Fix variable label type

metasyn/file.py

@@ -73,6 +73,21 @@ def to_dict(self) -> dict[str, Any]:
             "file_name": self.file_name,
         }
 
+    @classmethod
+    def from_dict(cls, interface_dict: dict):
+        """Create a file interface from the stored dictionary.
+
+        Parameters
+        ----------
+        interface_dict:
+            Dictionary containing the information to create the interface.
+        """
+        if interface_dict["file_interface_name"] != cls.format:
+            raise ValueError(f"Using wrong file interface, {interface_dict['file_interface_name']}"
+                             f" instead of {cls.format}.")
+        return cls(metadata=interface_dict["format_metadata"],
+                   file_name=interface_dict["file_name"])
+
     @abstractmethod
     def _write_file(self, df: pl.DataFrame, fp: Union[Path, str]):
         """Write synthetic file, to be implemented by file interface implementations.
@@ -185,6 +200,16 @@ class ReadStatInterface(BaseFileInterface, ABC):
 
     interface = "unknown"
 
+    def __init__(self, metadata, file_name):
+        super().__init__(metadata, file_name)
+        if "variable_value_labels" in metadata:
+            for col, data_dict in metadata["variable_value_labels"].items():
+                if not isinstance(data_dict, dict):
+                    continue
+                self.metadata["variable_value_labels"][col] = list(
+                    self.metadata["variable_value_labels"][col].items())
+
+
     def read_dataset(self, fp: Union[Path, str]):
         """Read the dataset without the metadata."""
         df, _ = SavFileInterface._get_df_metadata(fp)
@@ -284,7 +309,7 @@ def _write_file(self, df: pl.DataFrame, out_fp: Union[Path, str]):
             for col, mapping in self.metadata["variable_value_labels"].items():
                 if col not in df.columns or df[col].null_count() == len(df):
                     continue
-                reverse_mapping = {v: k for k, v in mapping.items()}
+                reverse_mapping = {v: k for k, v in mapping}
                 try:
                     # Try to replace the values back to the original with a fast method
                     out_df = out_df.with_columns(pl.col(col).replace_strict(reverse_mapping))
@@ -293,6 +318,9 @@ def _write_file(self, df: pl.DataFrame, out_fp: Union[Path, str]):
                     new_series = [x if x not in reverse_mapping else reverse_mapping[x]
                                   for x in out_df[col]]
                     df.replace_column(df.columns.index(col), pl.Series(col, new_series))
+            metadata["variable_value_labels"] = {
+                col: {k: v for k, v in data}
+                for col, data in metadata["variable_value_labels"].items()}
         metadata["variable_format"] = self._get_format(out_df, df)
         if not orig_file_label.startswith(label_start):
             file_label = label_start + orig_file_label
@@ -336,6 +364,7 @@ class SavFileInterface(ReadStatInterface):
     extensions = [".sav", ".zsav"]
     interface = "sav"
 
+
     @classmethod
     def _convert_with_orig_format(cls, df, prs_metadata):
         """Read the dataset including the metadata."""
@@ -351,6 +380,12 @@ def _convert_with_orig_format(cls, df, prs_metadata):
                     # Fast cast to integer
                     df = df.with_columns(pl.col(col).cast(pl.Int64))
 
+            if df[col].dtype == pl.Null:
+                if col_format.startswith("F") and col_format.endswith(".0"):
+                    df = df.with_columns(pl.col(col).cast(pl.Int64))
+                elif col_format.startswith("F"):
+                    df = df.with_columns(pl.col(col).cast(pl.Float64))
+
             # If column has one type, then it should be converted from object dtype to that type.
             if df[col].dtype == pl.Object and df[col].null_count() != len(df):
                 types = df[col].drop_nulls().map_elements(type)
@@ -389,6 +424,10 @@ def _prep_df_for_writing(self, df):
 
                 # Workaround bug in pyreadstat 1.3.7, remove when fixed.
                 if len(df) == df[col].null_count() and df[col].dtype == pl.Object:
+                    if col_format.startswith("F") and col_format.endswith(".0"):
+                        df = df.with_columns(pl.col(col).cast(pl.Int64))
+                    elif col_format.startswith("F"):
+                        df = df.with_columns(pl.col(col).cast(pl.Float64))
                     df = df.with_columns(pl.col(col).cast(pl.String))
         return df
 
@@ -598,8 +637,7 @@ def file_interface_from_dict(file_format_dict: dict) -> BaseFileInterface:
     """
     for interface_name, interface in _AVAILABLE_FILE_INTERFACES.items():
         if file_format_dict["file_interface_name"] == interface_name:
-            return interface(metadata=file_format_dict["format_metadata"],
-                             file_name=file_format_dict["file_name"])
+            return interface.from_dict(file_format_dict)
     raise ValueError(f"Cannot find file interface with name '{interface_name}'.")
 
 def read_file(fp: Union[Path, str], name: Optional[str] = None,

tests/test_file.py

@@ -20,6 +20,7 @@
     read_sav,
     read_tsv,
 )
+from metasyn.metaframe import MetaFrame
 
 
 class BadFileInterface1(BaseFileInterface):
@@ -136,18 +137,21 @@ def test_sav_interface(filename, tmpdir):
 
 @mark.parametrize("series,return_series,variable_format,labels,dtype", [
     (pl.Series([None, 1.0, 2.0]), pl.Series([None, "NA", 2.0], dtype=pl.Object),"F4.3", {1.0: "NA"}, pl.Object),
-    (pl.Series([None, 1, 2]), pl.Series([None, "NA", 2], dtype=pl.Object), "F8.0", {1: "NA"}, pl.Object),
-    (pl.Series([None, 2, 3]), pl.Series([None, 2, 3]), "F8.0", {1: "NA"}, pl.Int64),
-    (pl.Series([None, 1.5, 3.5]), pl.Series([None, 1.5, 3.5]), "F6.3", {1: "NA"}, pl.Float64),
-    (pl.Series([1, 2, 3]), pl.Series(["a", "b", "c"]), "F4.0", {1: "a", 2: "b", 3: "c"}, pl.Categorical),
+    (pl.Series([None, 1, 2]), pl.Series([None, "NA", 2], dtype=pl.Object), "F8.0", {1.0: "NA"}, pl.Object),
+    (pl.Series([None, 2, 3]), pl.Series([None, 2, 3]), "F8.0", {1.0: "NA"}, pl.Int64),
+    (pl.Series([None, 1.5, 3.5]), pl.Series([None, 1.5, 3.5]), "F6.3", {1.0: "NA"}, pl.Float64),
+    (pl.Series([1, 2, 3]), pl.Series(["a", "b", "c"]), "F4.0", {1.0: "a", 2.0: "b", 3.0: "c"}, pl.Categorical),
+    (pl.Series([None, None, None], dtype=pl.Int64), pl.Series([None, None, None], dtype=pl.Int64), "F6.0", {}, pl.Int64),
+    (pl.Series([None, None, None], dtype=pl.Float64), pl.Series([None, None, None], dtype=pl.Float64), "F6.2", {}, pl.Float64),
 ])
 def test_sav_dtype_file(series, return_series, variable_format, labels, dtype, tmpdir):
     test_file = tmpdir / "test.sav"
     df = pl.DataFrame({"test": series})
+    var_labels = {} if len(labels) == 0 else {"test": labels}
     metadata = {
         "column_labels": ["test_label"], "compress": False,
         "variable_format": {"test": variable_format}, "variable_display_width": {"test": 8},
-        "variable_value_labels": {"test": labels},
+        "variable_value_labels": var_labels,
         "variable_measure": {"test": "scale"},
         "apply_value_formats": True,
     }
@@ -156,12 +160,21 @@ def test_sav_dtype_file(series, return_series, variable_format, labels, dtype, t
     new_df, ff2 = SavFileInterface.read_file(test_file)
     ff2.metadata.pop("file_label")
     assert ff.metadata == ff2.metadata
+    new_label_keys = list(k for k, _ in ff2.metadata["variable_value_labels"].get("test", {}))
+    assert new_label_keys == list(labels)
+    assert all(type(new) is type(old) for new, old in zip(new_label_keys, list(labels)))
     assert all(x == y for x, y in zip(new_df["test"], return_series)), (new_df["test"], return_series)
     assert all(type(x) is type(y) for x, y in zip(new_df["test"], return_series)), [type(y) for y in new_df["test"]]
     assert new_df["test"].dtype == dtype
-    df_na, _ = SavFileInterface.read_file(test_file, apply_value_formats=False)
+    df_na, ff = SavFileInterface.read_file(test_file, apply_value_formats=False)
     assert all(x == y for x, y in zip(df_na["test"], series)), (df_na["test"], series)
-
+    mf = MetaFrame.fit_dataframe(df_na, file_format=ff)
+    mf.save_json(tmpdir / "tmp.json")
+    mf_new = MetaFrame.load_json(tmpdir / "tmp.json")
+    print(mf_new.file_format)
+    new_labels = mf_new.file_format["format_metadata"]["variable_value_labels"].get("test", [])
+    assert new_labels == list([k, v] for k, v in labels.items())
+    assert [[type(k), type(v)] for k, v in new_labels] == list([type(k), type(v)] for k, v in labels.items())
 
 @mark.parametrize(
     "max_rows,chunk_size,length",