Support generators for --convert functions, closes #16

README.md

@@ -146,6 +146,33 @@ You can import additional modules using `--import`. This example shows how you c
 
 If your Python code spans more than one line it needs to include a `return` statement.
 
+You can also use Python generators in your `--convert` code, for example:
+
+    git-history file stats.db package-stats/stats.json \
+        --repo package-stats \
+        --convert '
+        data = json.loads(content)
+        for key, counts in data.items():
+            for date, count in counts.items():
+                yield {
+                    "package": key,
+                    "date": date,
+                    "count": count
+                }
+        ' --id package --id date
+
+This conversion function expects data that looks like this:
+
+```json
+{
+    "airtable-export": {
+        "2021-05-18": 66,
+        "2021-05-19": 60,
+        "2021-05-20": 87
+    }
+}
+```
+
 ## Development
 
 To contribute to this tool, first checkout the code. Then create a new virtual environment:

git_history/cli.py

@@ -165,7 +165,8 @@ def file(
             # Skip empty JSON files
             continue
 
-        items = fn(content)
+        # list() to resolve generators for repeated access later
+        items = list(fn(content))
 
         # Remove any --ignore columns
         if ignore:

tests/test_git_history.py

@@ -342,3 +342,56 @@ def test_csv_tsv(repo, tmpdir, file):
         "   PRIMARY KEY ([_item], [_version])\n"
         ");"
     )
+
+
+@pytest.mark.parametrize(
+    "convert,expected_rows",
+    (
+        (
+            "json.loads(content.upper())",
+            [
+                {"ITEM_ID": 1, "NAME": "GIN"},
+                {"ITEM_ID": 2, "NAME": "TONIC"},
+                {"ITEM_ID": 1, "NAME": "GIN"},
+                {"ITEM_ID": 2, "NAME": "TONIC 2"},
+                {"ITEM_ID": 3, "NAME": "RUM"},
+            ],
+        ),
+        # Generator
+        (
+            (
+                "data = json.loads(content)\n"
+                "for item in data:\n"
+                '    yield {"just_name": item["name"]}'
+            ),
+            [
+                {"just_name": "Gin"},
+                {"just_name": "Tonic"},
+                {"just_name": "Gin"},
+                {"just_name": "Tonic 2"},
+                {"just_name": "Rum"},
+            ],
+        ),
+    ),
+)
+def test_convert(repo, tmpdir, convert, expected_rows):
+    runner = CliRunner()
+    db_path = str(tmpdir / "db.db")
+    with runner.isolated_filesystem():
+        result = runner.invoke(
+            cli,
+            [
+                "file",
+                db_path,
+                str(repo / "items.json"),
+                "--repo",
+                str(repo),
+                "--convert",
+                convert,
+            ],
+            catch_exceptions=False,
+        )
+    assert result.exit_code == 0
+    db = sqlite_utils.Database(db_path)
+    rows = [{k: v for k, v in r.items() if k != "_commit"} for r in db["items"].rows]
+    assert rows == expected_rows