fix string input in csv source2dicts

CentreForDigitalHumanities · Apr 11, 2024 · d35e888 · d35e888
1 parent 0372a6e
commit d35e888
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/ianalyzer_readers/readers/csv.py b/ianalyzer_readers/readers/csv.py
@@ -76,7 +76,7 @@ def source2dicts(self, source: Source) -> Iterable[Document]:
         if isinstance(source, str):
             filename = source
             metadata = {}
-        if isinstance(source, bytes):
+        elif isinstance(source, bytes):
             raise NotImplementedError()
         else:
             filename, metadata = source

diff --git a/tests/mock_csv_corpus.py b/tests/mock_csv_corpus.py
@@ -15,9 +15,7 @@ class TestCSVReader(CSVReader):
     def sources(self, **kwargs):
         for filename in os.listdir(self.data_directory):
             full_path = os.path.join(self.data_directory, filename)
-            yield full_path, {
-                'filename': filename
-            }
+            yield full_path
 
     fields = [
         Field(

diff --git a/tests/test_csvcorpus.py b/tests/test_csvcorpus.py
@@ -49,8 +49,20 @@ def test_csv():
     corpus = TestCSVReader()
 
     sources = list(corpus.sources())
-    assert len(sources) == 1 and sources[0][1] == {'filename': 'example.csv'}
+    assert len(sources) == 1
 
     docs = corpus.source2dicts(sources[0])
     for doc, target in zip(docs, target_documents):
         assert doc == target
+
+def test_csv_supported_source_types():
+    corpus = TestCSVReader()
+    source = next(corpus.sources())
+    assert isinstance(source, str)
+
+    # should work with a path as the source
+    list(corpus.source2dicts(source))
+
+    # should work with a path + metadata as the source
+    source_with_metadata = source, {}
+    list(corpus.source2dicts(source_with_metadata))