source: csv: Enable setting repo src_url using CSVSourceConfig

* Add test_key in CSVTest
* Add CSVSourceConfig

Author: sudharsana-kjl

CHANGELOG.md

@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   config for them.
 - shouldi example uses updated `MemoryOrchestrator.basic_config` method and
   includes more explanation in comments.
+- CSVSource allows for setting the Repo's `src_url` from a csv column
 ### Fixed
 - Docs get version from dffml.version.VERSION.
 - FileSource zipfiles are wrapped with TextIOWrapper because CSVSource expects

dffml/source/csv.py

@@ -5,14 +5,23 @@
 """
 import csv
 import ast
+from typing import NamedTuple, Dict
 
 from ..repo import Repo
 from .memory import MemorySource
-from .file import FileSource
+from .file import FileSource, FileSourceConfig
+from ..util.cli.arg import Arg
 
 csv.register_dialect("strip", skipinitialspace=True)
 
 
+class CSVSourceConfig(FileSourceConfig, NamedTuple):
+    filename: str
+    label: str = "unlabeled"
+    readonly: bool = False
+    key: str = None
+
+
 class CSVSource(FileSource, MemorySource):
     """
     Uses a CSV file as the source of repo feature data
@@ -21,6 +30,30 @@ class CSVSource(FileSource, MemorySource):
     # Headers we've added to track data other than feature data for a repo
     CSV_HEADERS = ["prediction", "confidence", "classification"]
 
+    @classmethod
+    def args(cls, args, *above) -> Dict[str, Arg]:
+        cls.config_set(args, above, "filename", Arg())
+        cls.config_set(
+            args,
+            above,
+            "readonly",
+            Arg(type=bool, action="store_true", default=False),
+        )
+        cls.config_set(
+            args, above, "label", Arg(type=str, default="unlabeled")
+        )
+        cls.config_set(args, above, "key", Arg(type=str, default=None))
+        return args
+
+    @classmethod
+    def config(cls, config, *above):
+        return CSVSourceConfig(
+            filename=cls.config_get(config, above, "filename"),
+            readonly=cls.config_get(config, above, "readonly"),
+            label=cls.config_get(config, above, "label"),
+            key=cls.config_get(config, above, "key"),
+        )
+
     async def load_fd(self, fd):
         """
         Parses a CSV stream into Repo instances
@@ -45,6 +78,11 @@ async def load_fd(self, fd):
                         repo_data["features"][key] = ast.literal_eval(value)
                     except (SyntaxError, ValueError):
                         repo_data["features"][key] = value
+                if self.config.key is not None and self.config.key == key:
+                    src_url = value
+                if self.config.key is None:
+                    src_url = str(i)
+            i += 1
             # Correct types and structure of repo data from csv_meta
             if "classification" in csv_meta:
                 repo_data.update(
@@ -59,9 +97,7 @@ async def load_fd(self, fd):
                         }
                     }
                 )
-            # Create the repo with the source URL being the row index
-            repo = Repo(str(i), data=repo_data)
-            i += 1
+            repo = Repo(src_url, data=repo_data)
             self.mem[repo.src_url] = repo
         self.logger.debug("%r loaded %d records", self, len(self.mem))
 

tests/source/test_csv.py

@@ -1,19 +1,63 @@
 # SPDX-License-Identifier: MIT
 # Copyright (c) 2019 Intel Corporation
 import unittest
+import tempfile
 
 from dffml.source.file import FileSourceConfig
-from dffml.source.csv import CSVSource
+from dffml.source.csv import CSVSource, CSVSourceConfig
 from dffml.util.testing.source import FileSourceTest
 from dffml.util.asynctestcase import AsyncTestCase
+from dffml.repo import Repo
+from dffml.util.cli.arg import parse_unknown
 
 
 class TestCSVSource(FileSourceTest, AsyncTestCase):
     async def setUpSource(self):
-        return CSVSource(FileSourceConfig(filename=self.testfile))
+        return CSVSource(CSVSourceConfig(filename=self.testfile))
 
     @unittest.skip("Labels not implemented yet for CSV files")
     async def test_label(self):
         """
         Labels not implemented yet for CSV files
         """
+
+    def test_config_readonly_default(self):
+        config = CSVSource.config(
+            parse_unknown("--source-file-filename", "feedface")
+        )
+        self.assertEqual(config.filename, "feedface")
+        self.assertEqual(config.label, "unlabeled")
+        self.assertEqual(config.key, None)
+        self.assertFalse(config.readonly)
+
+    def test_config_readonly_set(self):
+        config = CSVSource.config(
+            parse_unknown(
+                "--source-file-filename",
+                "feedface",
+                "--source-file-label",
+                "default-label",
+                "--source-file-key",
+                "SourceURLColumn",
+                "--source-file-readonly",
+            )
+        )
+        self.assertEqual(config.filename, "feedface")
+        self.assertEqual(config.label, "default-label")
+        self.assertEqual(config.key, "SourceURLColumn")
+        self.assertTrue(config.readonly)
+
+    async def test_key(self):
+        with tempfile.NamedTemporaryFile() as fileobj:
+            fileobj.write(b"KeyHeader,ValueColumn\n")
+            fileobj.write(b"a,42\n")
+            fileobj.write(b"b,420\n")
+            fileobj.seek(0)
+            async with CSVSource(
+                CSVSourceConfig(filename=fileobj.name, key="KeyHeader")
+            ) as source:
+                async with source() as sctx:
+                    repo_a = await sctx.repo("a")
+                    repo_b = await sctx.repo("b")
+                    self.assertEqual(repo_a.data.features["ValueColumn"], 42)
+                    self.assertEqual(repo_b.data.features["ValueColumn"], 420)