ENH: Allow pulling of partial results (#46)

Author: vegard-solum-4ss

fourinsight/engineroom/utils/_core.py

@@ -436,42 +436,51 @@ def append(self, dataframe):
             self.new_row(row_i)
             self.collect(**result_i)
 
-    def pull(self, raise_on_missing=True):
+    def pull(self, raise_on_missing=True, strict=True):
         """
         Pull results from source. Remote source overwrites existing values.
 
         Parameters
         ----------
         raise_on_missing : bool
             Raise exception if results can not be pulled from source.
+        strict : bool
+            Whether to be strict with respect to headers in the source or not. Setting
+            `strict=True` (default) will require that the source has the exact same
+            headers as the `ResultCollector`. Setting `strict=False` will allow pulling
+            of partial results (i.e., headers that does not have results in source,
+            will be populated with `None` values).
         """
 
         self._handler.pull(raise_on_missing=raise_on_missing)
         if not self._handler.getvalue():
             return
 
         self._handler.seek(0)
-        df = pd.read_csv(
+        df_source = pd.read_csv(
             self._handler, index_col=0, parse_dates=True, dtype=self._headers
         )
 
-        if not (set(df.columns) == set(self._headers.keys())):
+        if strict and set(df_source.columns) != set(self._headers.keys()):
             raise ValueError("Header is not valid.")
 
         if (
-            not df.index.empty
+            not df_source.index.empty
             and (self._indexing_mode == "auto")
-            and not (df.index.dtype == "int64")
+            and not (df_source.index.dtype == "int64")
         ):
             raise ValueError("Index dtype must be 'int64'.")
         elif (
-            not df.index.empty
+            not df_source.index.empty
             and (self._indexing_mode == "timestamp")
-            and not (isinstance(df.index, pd.DatetimeIndex))
+            and not (isinstance(df_source.index, pd.DatetimeIndex))
         ):
             raise ValueError("Index must be 'DatetimeIndex'.")
 
-        self._dataframe = df
+        columns_missing = self.dataframe.columns.difference(df_source.columns)
+        df_source[columns_missing] = None
+
+        self._dataframe = df_source[self._headers.keys()].astype(self._headers)
 
     def push(self):
         """

tests/test_core.py

@@ -667,6 +667,65 @@ def test_pull_non_existing_raises(self):
         with pytest.raises(FileNotFoundError):
             results.pull(raise_on_missing=True)
 
+    def test_pull_strict(self, tmp_path):
+        handler = LocalFileHandler(tmp_path / "results.csv")
+
+        headers_a = {"a": float, "b": str, "c": int}
+        results_a = ResultCollector(headers_a, handler=handler)
+        results_a.new_row()
+        results_a.collect(a=1.2, b="foo", c=3)
+        results_a.new_row()
+        results_a.collect(a=4.5, b="bar", c=6)
+        results_a.push()
+
+        headers_b = {"b": str, "c": int, "d": float}  # partially different headers
+        results_b = ResultCollector(headers_b, handler=handler)
+        results_b.pull(strict=False)
+
+        results_b.push()
+        results_a.pull(strict=False)
+
+        df_out_a = results_a.dataframe
+        df_out_b = results_b.dataframe
+
+        df_expect_a = pd.DataFrame(
+            data={
+                "a": [None, None],
+                "b": ["foo", "bar"],
+                "c": [3, 6],
+            },
+            index=pd.Index([0, 1], dtype="int64"),
+        ).astype({"a": "float64", "b": "string", "c": "Int64"})
+
+        df_expect_b = pd.DataFrame(
+            data={
+                "b": ["foo", "bar"],
+                "c": [3, 6],
+                "d": [None, None],
+            },
+            index=pd.Index([0, 1], dtype="int64"),
+        ).astype({"b": "string", "c": "Int64", "d": "float64"})
+
+        # Remove `check_index_type=False` when issue with index type in `pull` is fixed
+        pd.testing.assert_frame_equal(df_out_a, df_expect_a, check_index_type=False)
+        pd.testing.assert_frame_equal(df_out_b, df_expect_b, check_index_type=False)
+
+    def test_pull_strict_raises(self, tmp_path):
+        handler = LocalFileHandler(tmp_path / "results.csv")
+
+        headers_a = {"a": float, "b": str, "c": int}
+        results_a = ResultCollector(headers_a, handler=handler)
+        results_a.new_row()
+        results_a.collect(a=1.2, b="foo", c=3)
+        results_a.new_row()
+        results_a.collect(a=4.5, b="bar", c=6)
+        results_a.push()
+
+        headers_b = {"b": str, "c": int, "d": float}  # different headers
+        results_b = ResultCollector(headers_b, handler=handler)
+        with pytest.raises(ValueError):
+            results_b.pull()
+
     def test_push_auto(self, tmp_path):
         handler = LocalFileHandler(tmp_path / "results.csv")
         headers = {"a": float, "b": str, "c": float, "d": str}