support reading parquet files

setup.py

@@ -34,6 +34,7 @@
         "pandas",
         "pyyaml>=5.1",
         "vivarium",
+        "pyarrow",
     ]
 
     interactive_requirements = [

src/pseudopeople/entity_types.py

@@ -58,7 +58,12 @@ def __call__(
         randomness_stream: RandomnessStream,
         additional_key: Any,
     ) -> pd.Series:
-        column = column.copy()
+        # TODO: this is a temporary hack to account for all string columns having been made categorical
+        #  We should record expected output dtype in the columns data structure
+        if column.dtype.name == "category":
+            column = column.astype(str)
+        else:
+            column = column.copy()
         noise_level = configuration.row_noise_level
         to_noise_idx = get_index_to_noise(
             column, noise_level, randomness_stream, f"{self.name}_{additional_key}"

src/pseudopeople/interface.py

@@ -30,12 +30,27 @@ def _generate_form(
         Noised form data
     """
     configuration_tree = get_configuration(configuration)
+    if isinstance(source, str):
+        source = Path(source)
     if isinstance(source, pd.DataFrame):
         data = source
+    elif isinstance(source, Path):
+        if source.suffix == ".hdf":
+            data = pd.read_hdf(source)
+            if not isinstance(data, pd.DataFrame):
+                raise TypeError(f"File located at {source} must contain a pandas DataFrame.")
+        elif source.suffix == ".parquet":
+            data = pd.read_parquet(source)
+        else:
+            raise ValueError(
+                "Source path must either be a .hdf or a .parquet file. Provided "
+                f"{source.suffix}"
+            )
     else:
-        data = pd.read_hdf(source)
-        if not isinstance(data, pd.DataFrame):
-            raise TypeError(f"File located at {source} must contain a pandas DataFrame.")
+        raise TypeError(
+            f"Source {source} must be either a pandas DataFrame or a path to a "
+            "file containing a pandas DataFrame."
+        )
     return noise_form(form, data, configuration_tree, seed)
 
 

tests/unit/test_noise_form.py

@@ -170,7 +170,7 @@ def test_correct_forms_are_used(func, form, mocker):
         pytest.skip(reason=f"TODO: implement function for {form.value} form")
     mock = mocker.patch("pseudopeople.interface.noise_form")
     mocker.patch("pseudopeople.interface.pd.read_hdf", return_value=pd.DataFrame())
-    _ = func("dummy/path")
+    _ = func("dummy/path.hdf")
 
     assert mock.call_args[0][0] == form