Merge pull request #116 from Point72/tkp/etl

timkpaine · web-flow · commit 99b805dee99d · 2025-10-31T16:34:24.000-04:00
Add end-to-end etl example
diff --git a/.gitignore b/.gitignore
@@ -157,3 +157,11 @@ ccflow/labextension
 
 # Rust
 target
+
+# Examples
+outputs
+raw.html
+extracted.csv
+etl.db
+lobsters.html
+lobsters.csv
diff --git a/ccflow/callable.py b/ccflow/callable.py
@@ -15,7 +15,7 @@
 import logging
 from functools import lru_cache, wraps
 from inspect import Signature, isclass, signature
-from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, get_args, get_origin
 
 from pydantic import BaseModel as PydanticBaseModel, ConfigDict, Field, InstanceOf, PrivateAttr, TypeAdapter, field_validator, model_validator
 from typing_extensions import override
@@ -217,7 +217,9 @@ def get_evaluation_context(model: CallableModelType, context: ContextType, as_di
         def wrapper(model, context=Signature.empty, *, _options: Optional[FlowOptions] = None, **kwargs):
             if not isinstance(model, CallableModel):
                 raise TypeError(f"Can only decorate methods on CallableModels (not {type(model)}) with the flow decorator.")
-            if not isclass(model.context_type) or not issubclass(model.context_type, ContextBase):
+            if (not isclass(model.context_type) or not issubclass(model.context_type, ContextBase)) and not (
+                get_origin(model.context_type) is Union and type(None) in get_args(model.context_type)
+            ):
                 raise TypeError(f"Context type {model.context_type} must be a subclass of ContextBase")
             if not isclass(model.result_type) or not issubclass(model.result_type, ResultBase):
                 raise TypeError(f"Result type {model.result_type} must be a subclass of ResultBase")
@@ -237,7 +239,11 @@ def wrapper(model, context=Signature.empty, *, _options: Optional[FlowOptions] =
 
             # Type coercion on input. We do this here (rather than relying on ModelEvaluationContext) as it produces a nicer traceback/error message
             if not isinstance(context, model.context_type):
-                context = model.context_type.model_validate(context)
+                if get_origin(model.context_type) is Union and type(None) in get_args(model.context_type):
+                    model_context_type = [t for t in get_args(model.context_type) if t is not type(None)][0]
+                else:
+                    model_context_type = model.context_type
+                context = model_context_type.model_validate(context)
 
             if fn != getattr(model.__class__, fn.__name__).__wrapped__:
                 # This happens when super().__call__ is used when implementing a CallableModel that derives from another one.
@@ -385,7 +391,8 @@ class ModelEvaluationContext(
     fn: str = Field("__call__", strict=True)
     options: Dict[str, Any] = Field(default_factory=dict)
     model: InstanceOf[_CallableModel]
-    context: InstanceOf[ContextBase]
+    context: Union[InstanceOf[ContextBase], None]
+
     # Using InstanceOf instead of the actual type will limit Pydantic's validation of the field to instance checking
     # Otherwise, the validation will re-run fully despite the models already being validated on construction
     # TODO: Make the instance check compatible with the generic types instead of the base type
@@ -492,9 +499,15 @@ def context_type(self) -> Type[ContextType]:
         typ = _cached_signature(self.__class__.__call__).parameters["context"].annotation
         if typ is Signature.empty:
             raise TypeError("Must either define a type annotation for context on __call__ or implement 'context_type'")
-        if not issubclass(typ, ContextBase):
-            raise TypeError(f"Context type declared in signature of __call__ must be a subclass of ContextBase. Received {typ}.")
 
+        # If optional type, extract inner type
+        if get_origin(typ) is Optional or (get_origin(typ) is Union and type(None) in get_args(typ)):
+            typ_to_check = [t for t in get_args(typ) if t is not type(None)][0]
+        else:
+            typ_to_check = typ
+        # Ensure subclass of ContextBase
+        if not issubclass(typ_to_check, ContextBase):
+            raise TypeError(f"Context type declared in signature of __call__ must be a subclass of ContextBase. Received {typ_to_check}.")
         return typ
 
     @property
diff --git a/ccflow/examples/etl/__init__.py b/ccflow/examples/etl/__init__.py
@@ -0,0 +1,25 @@
+from pathlib import Path
+from typing import List, Optional
+
+from ccflow import RootModelRegistry, load_config as load_config_base
+
+__all__ = ("load_config",)
+
+
+def load_config(
+    config_dir: str = "",
+    config_name: str = "",
+    overrides: Optional[List[str]] = None,
+    *,
+    overwrite: bool = True,
+    basepath: str = "",
+) -> RootModelRegistry:
+    return load_config_base(
+        root_config_dir=str(Path(__file__).resolve().parent / "config"),
+        root_config_name="base",
+        config_dir=config_dir,
+        config_name=config_name,
+        overrides=overrides,
+        overwrite=overwrite,
+        basepath=basepath,
+    )
diff --git a/ccflow/examples/etl/__main__.py b/ccflow/examples/etl/__main__.py
@@ -0,0 +1,39 @@
+import hydra
+
+from ccflow.utils.hydra import cfg_run
+
+__all__ = ("main",)
+
+
+@hydra.main(config_path="config", config_name="base", version_base=None)
+def main(cfg):
+    cfg_run(cfg)
+
+
+# Extract step:
+# python -m ccflow.examples.etl +callable=extract +context=[]
+# Change url, as example of context override:
+# python -m ccflow.examples.etl +callable=extract +context=["http://lobste.rs"]
+# Change file name, as example of callable override:
+# python -m ccflow.examples.etl +callable=extract +context=["http://lobste.rs"] ++extract.publisher.name=lobsters
+
+# Transform step:
+# python -m ccflow.examples.etl +callable=transform +context=[]
+# python -m ccflow.examples.etl +callable=transform +context=[] ++transform.model.file=lobsters.html ++transform.publisher.name=lobsters
+
+# Load step:
+# python -m ccflow.examples.etl +callable=load +context=[]
+# python -m ccflow.examples.etl +callable=load +context=[] ++load.file=lobsters.csv ++load.db_file=":memory:"
+
+# View SQLite DB:
+# sqlite3 etl.db
+# .tables
+# select * from links;
+# .quit
+
+# [project.scripts]
+# etl = "ccflow.examples.etl:main"
+# etl-explain = "ccflow.examples.etl:explain"
+
+if __name__ == "__main__":
+    main()
diff --git a/ccflow/examples/etl/config/__init__.py b/ccflow/examples/etl/config/__init__.py
diff --git a/ccflow/examples/etl/config/base.yaml b/ccflow/examples/etl/config/base.yaml
@@ -0,0 +1,32 @@
+extract:
+  _target_: ccflow.PublisherModel
+  model:
+    _target_: ccflow.examples.etl.models.RestModel
+  publisher:
+    _target_: ccflow.publishers.GenericFilePublisher
+    name: raw
+    suffix: .html
+  field: value
+
+transform:
+  _target_: ccflow.PublisherModel
+  model:
+    _target_: ccflow.examples.etl.models.LinksModel
+    file: ${extract.publisher.name}${extract.publisher.suffix}
+  publisher:
+    _target_: ccflow.publishers.GenericFilePublisher
+    name: extracted
+    suffix: .csv
+  field: value
+
+load:
+  _target_: ccflow.examples.etl.models.DBModel
+  file: ${transform.publisher.name}${transform.publisher.suffix}
+  db_file: etl.db
+  table: links
+
+# Alternative multi-file approach
+# defaults:
+#     - extract: rest
+#     - transform: links
+#     - load: db
diff --git a/ccflow/examples/etl/config/extract/rest.yaml b/ccflow/examples/etl/config/extract/rest.yaml
@@ -0,0 +1,8 @@
+_target_: ccflow.PublisherModel
+model:
+  _target_: ccflow.examples.etl.models.RestModel
+publisher:
+  _target_: ccflow.publishers.GenericFilePublisher
+  name: raw
+  suffix: .html
+field: value
diff --git a/ccflow/examples/etl/config/load/db.yaml b/ccflow/examples/etl/config/load/db.yaml
@@ -0,0 +1,4 @@
+_target_: ccflow.examples.etl.models.DBModel
+file: ${transform.publisher.name}${transform.publisher.suffix}
+db_file: etl.db
+table: links
diff --git a/ccflow/examples/etl/config/transform/links.yaml b/ccflow/examples/etl/config/transform/links.yaml
@@ -0,0 +1,9 @@
+_target_: ccflow.PublisherModel
+model:
+  _target_: ccflow.examples.etl.models.LinksModel
+  file: ${extract.publisher.name}${extract.publisher.suffix}
+publisher:
+  _target_: ccflow.publishers.GenericFilePublisher
+  name: extracted
+  suffix: .csv
+field: value
diff --git a/ccflow/examples/etl/explain.py b/ccflow/examples/etl/explain.py
@@ -0,0 +1,13 @@
+from ccflow.utils.hydra import cfg_explain_cli
+
+from .__main__ import main
+
+__all__ = ("explain",)
+
+
+def explain():
+    cfg_explain_cli(config_path="config", config_name="base", hydra_main=main)
+
+
+if __name__ == "__main__":
+    explain()
diff --git a/ccflow/examples/etl/models.py b/ccflow/examples/etl/models.py
@@ -0,0 +1,76 @@
+import sqlite3
+from csv import DictReader, DictWriter
+from io import StringIO
+from typing import Optional
+
+from bs4 import BeautifulSoup
+from httpx import Client
+from pydantic import Field
+
+from ccflow import CallableModel, ContextBase, Flow, GenericResult, NullContext
+
+__all__ = ("RestModel", "LinksModel", "DBModel", "SiteContext")
+
+
+class SiteContext(ContextBase):
+    """An example of a context object, passed into and between callable models from the command line."""
+
+    site: str = Field(default="https://en.wikipedia.org/wiki/Main_Page")
+
+
+class RestModel(CallableModel):
+    """Example callable model that fetches a URL and returns the HTML content."""
+
+    @Flow.call
+    def __call__(self, context: Optional[SiteContext] = None) -> GenericResult[str]:
+        context = context or SiteContext()
+        resp = Client().get(context.site, headers={"User-Agent": "Safari/537.36"}, follow_redirects=True)
+        resp.raise_for_status()
+
+        return GenericResult[str](value=resp.text)
+
+
+class LinksModel(CallableModel):
+    """Example callable model that transforms HTML content into CSV of links."""
+
+    file: str
+
+    @Flow.call
+    def __call__(self, context: Optional[NullContext] = None) -> GenericResult[str]:
+        context = context or NullContext()
+
+        with open(self.file, "r") as f:
+            html = f.read()
+
+        # Use beautifulsoup to convert links into csv of name, url
+        soup = BeautifulSoup(html, "html.parser")
+        links = [{"name": a.text, "url": href} for a in soup.find_all("a", href=True) if (href := a["href"]).startswith("http")]
+
+        io = StringIO()
+        writer = DictWriter(io, fieldnames=["name", "url"])
+        writer.writeheader()
+        writer.writerows(links)
+        output = io.getvalue()
+        return GenericResult[str](value=output)
+
+
+class DBModel(CallableModel):
+    """Example callable model that loads CSV data into a SQLite database."""
+
+    file: str
+    db_file: str = Field(default="etl.db")
+    table: str = Field(default="links")
+
+    @Flow.call
+    def __call__(self, context: Optional[NullContext] = None) -> GenericResult[str]:
+        context = context or NullContext()
+
+        conn = sqlite3.connect(self.db_file)
+        cursor = conn.cursor()
+        cursor.execute(f"CREATE TABLE IF NOT EXISTS {self.table} (name TEXT, url TEXT)")
+        with open(self.file, "r") as f:
+            reader = DictReader(f)
+            for row in reader:
+                cursor.execute(f"INSERT INTO {self.table} (name, url) VALUES (?, ?)", (row["name"], row["url"]))
+        conn.commit()
+        return GenericResult[str](value="Data loaded into database")
diff --git a/ccflow/result/__init__.py b/ccflow/result/__init__.py
@@ -1,2 +1,3 @@
 from .dict import *
 from .generic import *
+from .list import *
diff --git a/ccflow/result/list.py b/ccflow/result/list.py
@@ -0,0 +1,12 @@
+from typing import Generic, List, TypeVar
+
+from ..base import ResultBase
+
+__all__ = ("ListResult",)
+
+
+V = TypeVar("V")
+
+
+class ListResult(ResultBase, Generic[V]):
+    value: List[V]
diff --git a/ccflow/tests/examples/test_etl.py b/ccflow/tests/examples/test_etl.py
@@ -0,0 +1,50 @@
+from tempfile import NamedTemporaryFile
+from unittest.mock import patch
+
+from ccflow.examples.etl.__main__ import main
+from ccflow.examples.etl.explain import explain
+from ccflow.examples.etl.models import DBModel, LinksModel, RestModel, SiteContext
+
+
+class TestEtl:
+    def test_rest_model(self):
+        rest = RestModel()
+        context = SiteContext(site="https://en.wikipedia.org/wiki/Main_Page")
+        result = rest(context)
+        assert result.value is not None
+        assert "Wikipedia" in result.value
+
+    def test_links_model(self):
+        with NamedTemporaryFile(suffix=".html") as file:
+            file.write(b"""
+            <html>
+                <body>
+                    <a href="https://example.com/page1">Page 1</a>
+                    <a href="https://example.com/page2">Page 2</a>
+                </body>
+            </html>
+            """)
+            file.flush()
+            links = LinksModel(file=file.name)
+            result = links()
+            assert result.value is not None
+            assert "name,url" in result.value  # Check for CSV header
+
+    def test_db_model(self):
+        with NamedTemporaryFile(suffix=".csv", mode="w+", delete=False) as file:
+            file.write("name,url\nPage 1,https://example.com/page1\nPage 2,https://example.com/page2\n")
+            file.flush()
+            db = DBModel(file=file.name, db_file=":memory:", table="links")
+            result = db()
+            assert result.value == "Data loaded into database"
+
+    def test_cli(self):
+        with patch("ccflow.examples.etl.__main__.cfg_run") as mock_cfg_run:
+            with patch("sys.argv", ["etl", "+callable=extract", "+context=[]"]):
+                main()
+                mock_cfg_run.assert_called_once()
+
+    def test_explain(self):
+        with patch("ccflow.examples.etl.explain.cfg_explain_cli") as mock_cfg_explain:
+            explain()
+            mock_cfg_explain.assert_called_once()
diff --git a/ccflow/tests/test_callable.py b/ccflow/tests/test_callable.py
@@ -1,4 +1,4 @@
-from typing import Generic, List, TypeVar
+from typing import Generic, List, Optional, TypeVar
 from unittest import TestCase
 
 from pydantic import ValidationError
@@ -47,6 +47,13 @@ def __call__(self, context: MyContext) -> MyResult:
         return MyResult(x=self.i, y=context.a)
 
 
+class MyCallableOptionalContext(CallableModel):
+    @Flow.call
+    def __call__(self, context: Optional[MyContext] = None) -> MyResult:
+        context = context or MyContext(a="default")
+        return MyResult(x=1, y=context.a)
+
+
 class MyCallableChild(MyCallable):
     pass
 
@@ -215,6 +222,13 @@ def test_signature(self):
         self.assertRaises(TypeError, m, context, a="foo")
         self.assertRaises(TypeError, m, context=context, a="foo")
 
+    def test_signature_optional_context(self):
+        m = MyCallableOptionalContext()
+        context = MyContext(a="foo")
+        target = m(context)
+        self.assertEqual(m(context=context), target)
+        self.assertEqual(m().y, "default")
+
     def test_inheritance(self):
         m = MyCallableChild(i=5)
         self.assertEqual(m(MyContext(a="foo")), MyResult(x=5, y="foo"))
diff --git a/ccflow/tests/utils/test_hydra.py b/ccflow/tests/utils/test_hydra.py
diff --git a/ccflow/utils/hydra.py b/ccflow/utils/hydra.py
diff --git a/docs/img/wiki/etl/explain1.png b/docs/img/wiki/etl/explain1.png
diff --git a/docs/img/wiki/etl/explain2.png b/docs/img/wiki/etl/explain2.png
diff --git a/docs/wiki/ETL.md b/docs/wiki/ETL.md
diff --git a/pyproject.toml b/pyproject.toml
