Repository: intake/intake
Pull request #756 (2023-08-30 15:46:11)

Fix crash in CSVSource._get_schema when `skipfooter` is set in csv_kwargs
Fixes issue #755 


Patch:
diff --git a/intake/source/csv.py b/intake/source/csv.py
index df3eaf7..e446cfe 100644
--- a/intake/source/csv.py
+++ b/intake/source/csv.py
@@ -141,13 +141,30 @@ class CSVSource(base.DataSource, base.PatternMixin):
             self._schema = base.Schema(dtype=dtypes, shape=(None, len(dtypes)), npartitions=self._dask_df.npartitions, extra_metadata={})
             return self._schema
 
-        nrows = self._csv_kwargs.get("nrows")
-        self._csv_kwargs["nrows"] = 10
+        nrows = self._csv_kwargs.pop("nrows", None)
+        skipfooter = self._csv_kwargs.pop("skipfooter", None)
+        if skipfooter is not None:
+            try:
+                self._csv_kwargs["nrows"] = 10 + skipfooter
+            except TypeError:
+                raise TypeError("Non-numeric value passed for `skipfooter`.")
+        else:
+            self._csv_kwargs["nrows"] = 10
         df = self._get_partition(0)
+        # If `skipfooter` is not None, we need to check whether there are less
+        # than `10+skipfooter`` rows in `df`. If so, we need to re-read the
+        # DataFrame with `nrows` equal to the total number of rows in `df` minus
+        # `skipfooter`, to avoid that the footers interfere with detecting the
+        # right dtypes.
+        if skipfooter is not None and len(df) < 10 + skipfooter:
+            self._csv_kwargs["nrows"] = len(df) - skipfooter
+            df = self._get_partition(0)
         if nrows is None:
             del self._csv_kwargs["nrows"]
         else:
             self._csv_kwargs["nrows"] = nrows
+        if skipfooter is not None:
+            self._csv_kwargs["skipfooter"] = skipfooter
 
         dtypes = {col: str(dtype) for col, dtype in df.dtypes.items()}
         self._schema = base.Schema(dtype=dtypes, shape=(None, len(dtypes)), npartitions=len(self.files()), extra_metadata={})


Test patch:
diff --git a/intake/source/tests/footer_csvs/sample_fewfooters.csv b/intake/source/tests/footer_csvs/sample_fewfooters.csv
new file mode 100644
index 0000000..741e4db
--- /dev/null
+++ b/intake/source/tests/footer_csvs/sample_fewfooters.csv
@@ -0,0 +1,7 @@
+name,score,rank
+Alice,100.5,1
+Bob,50.3,2
+Charlie,25,3
+Eve,25,3
+1. This is footer #1
+2. This is another footer row, with a comma
diff --git a/intake/source/tests/footer_csvs/sample_manyfooters.csv b/intake/source/tests/footer_csvs/sample_manyfooters.csv
new file mode 100644
index 0000000..83d2458
--- /dev/null
+++ b/intake/source/tests/footer_csvs/sample_manyfooters.csv
@@ -0,0 +1,10 @@
+name,score,rank
+Alice,100.5,1
+Bob,50.3,2
+Charlie,25,3
+Eve,25,3
+1. This is footer #1
+2. This is another footer row, with a comma
+3. This a  third footer row. # With a hash
+4. And another one.
+5. And finally a last footer, to make sure we have more footers than data rows.
diff --git a/intake/source/tests/footer_csvs/sample_nofooters.csv b/intake/source/tests/footer_csvs/sample_nofooters.csv
new file mode 100644
index 0000000..d64b7f0
--- /dev/null
+++ b/intake/source/tests/footer_csvs/sample_nofooters.csv
@@ -0,0 +1,5 @@
+name,score,rank
+Alice,100.5,1
+Bob,50.3,2
+Charlie,25,3
+Eve,25,3
diff --git a/intake/source/tests/test_csv.py b/intake/source/tests/test_csv.py
index 530b83b..2ed78c1 100644
--- a/intake/source/tests/test_csv.py
+++ b/intake/source/tests/test_csv.py
@@ -66,6 +66,18 @@ def sample_pattern_datasource_with_cache(data_filenames):
     return csv.CSVSource(data_filenames["sample_pattern"], metadata=metadata)
 
 
+@pytest.fixture
+def footer_csv_dir():
+    return make_path_posix(os.path.join(os.path.dirname(__file__), "footer_csvs"))
+
+
+@pytest.fixture(params=[("sample_fewfooters.csv", 2), ("sample_manyfooters.csv", 5)])
+def sample_datasource_with_skipfooter(request, footer_csv_dir):
+    csv_filename, skipfooter = request.param
+    csv_path = posixpath.join(footer_csv_dir, csv_filename)
+    return csv.CSVSource(csv_path, csv_kwargs={"skipfooter": skipfooter})
+
+
 def test_csv_plugin():
     p = csv.CSVSource
     assert isinstance(p.version, str)
@@ -268,3 +280,13 @@ def test_pickle(sample1_datasource):
     df = sample1_clone.read()
 
     assert expected_df.equals(df)
+
+
+def test_skipfooter(sample_datasource_with_skipfooter, footer_csv_dir):
+    csv_nofooters = posixpath.join(footer_csv_dir, "sample_nofooters.csv")
+    expected_source = csv.CSVSource(csv_nofooters)
+    expected_schema = expected_source.discover()
+    assert expected_schema == sample_datasource_with_skipfooter.discover()
+    expected_df = expected_source.read()
+    df = sample_datasource_with_skipfooter.read()
+    assert expected_df.equals(df)