diff --git a/metasyn/file.py b/metasyn/file.py index d73bc413..80bed79e 100644 --- a/metasyn/file.py +++ b/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, diff --git a/tests/test_file.py b/tests/test_file.py index 3ee958c5..133ef983 100644 --- a/tests/test_file.py +++ b/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",