Update file updating logic

Author: gerlero

src/foamlib/_files/files.py

@@ -480,21 +480,6 @@ def getall(
             FoamFile.SubDict(self, keywords) if v is ... else deepcopy(v) for v in ret
         ]
 
-    def _normalize_and_validate_keywords(
-        self, keywords: str | tuple[str, ...] | None
-    ) -> tuple[str, ...]:
-        """Normalize keywords to tuple format and validate them."""
-        if keywords is None:
-            keywords = ()
-        elif not isinstance(keywords, tuple):
-            keywords = (keywords,)
-
-        if keywords and not isinstance(normalize(keywords[-1], bool_ok=False), str):
-            msg = f"Invalid keyword type: {keywords[-1]} (type {type(keywords[-1])})"
-            raise TypeError(msg)
-
-        return keywords
-
     def _write_header_if_needed(self, keywords: tuple[str, ...]) -> None:
         """Write FoamFile header if needed."""
         try:
@@ -517,10 +502,10 @@ def _write_header_if_needed(self, keywords: tuple[str, ...]) -> None:
     def _update_class_for_field_if_needed(
         self,
         keywords: tuple[str, ...],
-        data: DataLike | StandaloneDataLike | SubDictLike,
+        data: Data | StandaloneData | SubDict,
     ) -> None:
         """Update class field to appropriate field type if this is a field entry."""
-        if (
+        if isinstance(data, (float, int, np.ndarray)) and (
             keywords == ("internalField",)
             or (
                 len(keywords) == 3
@@ -532,26 +517,24 @@ def _update_class_for_field_if_needed(
                 )
             )
         ) and self.class_ == "dictionary":
-            try:
-                tensor_kind = _tensor_kind_for_field(data)  # ty: ignore[invalid-argument-type]
-            except ValueError:
-                pass
-            else:
-                self.class_ = "vol" + tensor_kind[0].upper() + tensor_kind[1:] + "Field"
+            tensor_kind = _tensor_kind_for_field(data)
+            self.class_ = "vol" + tensor_kind[0].upper() + tensor_kind[1:] + "Field"
 
     def _calculate_spacing(
         self,
         keywords: tuple[str, ...],
         start: int,
         end: int,
-        operation: Literal["put", "add"],
+        /,
+        *,
+        add: bool,
     ) -> tuple[bytes, bytes]:
         """Calculate before/after spacing for entry operations."""
         parsed = self._get_parsed(missing_ok=True)
 
         # For setitem operations, check if this is an update to an existing entry
         # and preserve existing spacing for sub-dictionary entries
-        if operation == "put":
+        if not add:
             is_update = keywords in parsed
             if is_update and len(keywords) > 1:
                 # For existing sub-dictionary entries, preserve existing formatting
@@ -588,105 +571,99 @@ def _calculate_spacing(
 
     def _perform_entry_operation(
         self,
-        keywords: str | tuple[str, ...] | None,
+        keywords: tuple[str, ...],
         data: DataLike | StandaloneDataLike | SubDictLike,
-        operation: Literal["put", "add"],
+        /,
+        *,
+        add: bool,
     ) -> None:
         """Shared method for performing entry operations (setitem and add)."""
-        keywords = self._normalize_and_validate_keywords(keywords)
+        if keywords:
+            keyword = normalize(keywords[-1], bool_ok=False)
+
+            if not isinstance(keyword, str):
+                msg = (
+                    f"Invalid keyword type: {keywords[-1]} (type {type(keywords[-1])})"
+                )
+                raise TypeError(msg)
+
+            if keyword != keywords[-1]:
+                msg = f"Invalid keyword: {keywords[-1]}"
+                raise ValueError(msg)
+
+        data = normalize(data, keywords=keywords)
+
+        indentation = b"    " * (len(keywords) - 1)
 
         with self:
             self._write_header_if_needed(keywords)
             self._update_class_for_field_if_needed(keywords, data)
 
             parsed = self._get_parsed(missing_ok=True)
-            start, end = parsed.entry_location(keywords, add=(operation == "add"))
-            before, after = self._calculate_spacing(keywords, start, end, operation)
-            self._process_data_entry(keywords, data, before, after, operation)
+            start, end = parsed.entry_location(keywords, add=add)
+            before, after = self._calculate_spacing(keywords, start, end, add=add)
 
-    def _process_data_entry(
-        self,
-        keywords: tuple[str, ...],
-        data: DataLike | StandaloneDataLike | SubDictLike,
-        before: bytes,
-        after: bytes,
-        operation: Literal["put", "add"],
-    ) -> None:
-        """Process and write a data entry using either put or add operation."""
-        parsed = self._get_parsed(missing_ok=True)
-        indentation = b"    " * (len(keywords) - 1)
+            try:
+                header = self["FoamFile"]
+                assert isinstance(header, FoamFile.SubDict)
+            except (KeyError, FileNotFoundError):
+                header = None
 
-        if isinstance(data, Mapping):
-            if not keywords:
-                msg = "Cannot set a mapping at the root level of a FoamFile\nUse update(), extend(), or merge() instead."
-                raise ValueError(msg)
+            if isinstance(data, Mapping):
+                if not keywords:
+                    msg = "Cannot set a mapping at the root level of a FoamFile\nUse update(), extend(), or merge() instead."
+                    raise ValueError(msg)
 
-            keyword = normalize(keywords[-1], bool_ok=False)
+                if keyword.startswith("#"):
+                    msg = f"Cannot set a directive as the keyword for a dictionary: {keyword}"
+                    raise ValueError(msg)
 
-            if not isinstance(keyword, str):
-                msg = (
-                    f"Invalid keyword type: {keywords[-1]} (type {type(keywords[-1])})"
-                )
-                raise TypeError(msg)
+                if add and keywords in parsed:
+                    raise KeyError(keywords)
 
-            if keyword.startswith("#"):
-                msg = (
-                    f"Cannot set a directive as the keyword for a dictionary: {keyword}"
+                empty_dict_content = (
+                    before
+                    + indentation
+                    + dumps(keyword)
+                    + b"\n"
+                    + indentation
+                    + b"{\n"
+                    + indentation
+                    + b"}"
+                    + after
                 )
-                raise ValueError(msg)
+                parsed.put(keywords, ..., empty_dict_content)
 
-            data = normalize(data, keywords=keywords)
-
-            content = (
-                before
-                + indentation
-                + dumps(keyword)
-                + b"\n"
-                + indentation
-                + b"{\n"
-                + indentation
-                + b"}"
-                + after
-            )
+                for k, v in data.items():
+                    self[(*keywords, k)] = v
 
-            if operation == "add" and keywords in parsed:
-                raise KeyError(keywords)
-
-            parsed.put(keywords, ..., content)
-            for k, v in data.items():  # ty: ignore[possibly-missing-attribute]
-                self[(*keywords, k)] = v
-
-        elif keywords:
-            header = self.get("FoamFile", None)
-            assert header is None or isinstance(header, FoamFile.SubDict)
-            val = dumps(data, keywords=keywords, header=header)
-
-            content = (
-                before
-                + indentation
-                + dumps(normalize(keywords[-1], bool_ok=False))
-                + ((b" " + val) if val else b"")
-                + (b";" if not keywords[-1].startswith("#") else b"")
-                + after
-            )
+            elif keywords:
+                val = dumps(data, keywords=keywords, header=header)
 
-            if operation == "put":
-                parsed.put(keywords, normalize(data, keywords=keywords), content)
-            else:  # operation == "add"
-                if keywords in parsed and not keywords[-1].startswith("#"):
-                    raise KeyError(keywords)
-                parsed.add(keywords, normalize(data, keywords=keywords), content)
+                content = (
+                    before
+                    + indentation
+                    + dumps(keyword)
+                    + ((b" " + val) if val else b"")
+                    + (b";" if not keywords[-1].startswith("#") else b"")
+                    + after
+                )
 
-        else:
-            if operation == "add" and () in parsed:
-                raise KeyError(())
+                if add:
+                    if keywords in parsed and not keywords[-1].startswith("#"):
+                        raise KeyError(keywords)
+
+                    parsed.add(keywords, data, content)
+                else:
+                    parsed.put(keywords, data, content)
 
-            header = self.get("FoamFile", None)
-            assert header is None or isinstance(header, FoamFile.SubDict)
+            else:
+                if add and () in parsed:
+                    raise KeyError(None)
 
-            content = before + dumps(data, keywords=(), header=header) + after
+                content = before + dumps(data, keywords=(), header=header) + after
 
-            parsed.put((), normalize(data, keywords=keywords), content)
+                parsed.put((), data, content)
 
     @overload
     def __setitem__(
@@ -709,15 +686,25 @@ def __setitem__(  # ty: ignore[invalid-method-override]
         keywords: str | tuple[str, ...] | None,
         data: DataLike | StandaloneDataLike | SubDictLike,
     ) -> None:
-        self._perform_entry_operation(keywords, data, "put")
+        if keywords is None:
+            keywords = ()
+        elif not isinstance(keywords, tuple):
+            keywords = (keywords,)
+
+        self._perform_entry_operation(keywords, data, add=False)
 
     @override
     def add(
         self,
         keywords: str | tuple[str, ...] | None,
         data: DataLike | StandaloneDataLike | SubDictLike,
     ) -> None:
-        self._perform_entry_operation(keywords, data, "add")
+        if keywords is None:
+            keywords = ()
+        elif not isinstance(keywords, tuple):
+            keywords = (keywords,)
+
+        self._perform_entry_operation(keywords, data, add=True)
 
     @with_default
     @override
@@ -969,11 +956,9 @@ def dumps(
         if header is None and ensure_header:
             class_ = "dictionary"
             if isinstance(file, Mapping) and "internalField" in file:
-                try:
-                    tensor_kind = _tensor_kind_for_field(file["internalField"])
-                except (ValueError, TypeError):
-                    pass
-                else:
+                internal_field = file["internalField"]
+                if isinstance(internal_field, (float, int, np.ndarray)):
+                    tensor_kind = _tensor_kind_for_field(internal_field)
                     class_ = "vol" + tensor_kind[0].upper() + tensor_kind[1:] + "Field"
 
             header = {"version": 2.0, "format": "ascii", "class": class_}