feat: Add support for loading context from TOML files including pyproject.toml

debug_toml.py

@@ -0,0 +1,7 @@
+import tomli
+from pathlib import Path
+
+toml_path = Path('tests/cli/resources/test_context.toml')
+with open(toml_path, 'rb') as f:
+    data = tomli.load(f)
+print('TOML data:', data)

hamilton/cli/__main__.py

@@ -81,7 +81,7 @@ class CliState:
     typer.Option(
         "--context",
         "-ctx",
-        help="Path to Driver context file [.json, .py]",
+        help="Path to Driver context file [.json, .py, .toml]",
         exists=True,
         dir_okay=False,
         readable=True,

hamilton/cli/logic.py

@@ -269,7 +269,6 @@ def visualize_diff(
 
 
 # TODO refactor ContextLoader to a class
-# TODO support loading from pyproject.toml
 def load_context(file_path: Path) -> dict:
     if not file_path.exists():
         raise FileNotFoundError(f"`{file_path}` doesn't exist.")
@@ -279,6 +278,8 @@ def load_context(file_path: Path) -> dict:
         context = _read_json_context(file_path)
     elif extension == ".py":
         context = _read_py_context(file_path)
+    elif extension in [".toml", ".tml"]:
+        context = _read_toml_context(file_path)
     else:
         raise ValueError(f"Received extension `{extension}` is unsupported.")
 
@@ -337,3 +338,43 @@ def _read_py_context(file_path: Path) -> dict:
         context[k] = getattr(module, k, None)
 
     return context
+
+
+def _read_toml_context(file_path: Path) -> dict:
+    """Read context from a TOML file. For pyproject.toml, looks for Hamilton configuration in [tool.hamilton] section."""
+    try:
+        import tomli  # Using tomli for compatibility with older Python versions
+    except ImportError:
+        # Provide a helpful error message if tomli is not available
+        raise ImportError(
+            "tomli is required to read TOML files. "
+            "Install it with `pip install tomli` or `pip install sf-hamilton[cli]` which includes TOML support."
+        )
+
+    with open(file_path, 'rb') as f:
+        data = tomli.load(f)
+
+    # First check if there's a [tool.hamilton] section in pyproject.toml
+    # This is where Hamilton-specific configuration would typically go
+    hamilton_config = data.get('tool', {}).get('hamilton', {})
+
+    # If we find Hamilton-specific config, use it as the context
+    if hamilton_config:
+        context = {
+            CONFIG_HEADER: hamilton_config.get('config', {}),
+            FINAL_VARS_HEADER: hamilton_config.get('final_vars', []),
+            INPUTS_HEADER: hamilton_config.get('inputs', {}),
+            OVERRIDES_HEADER: hamilton_config.get('overrides', {}),
+        }
+    else:
+        # Otherwise, check for top-level Hamilton context headers
+        context = {}
+        for k in [
+            CONFIG_HEADER,
+            FINAL_VARS_HEADER,
+            INPUTS_HEADER,
+            OVERRIDES_HEADER,
+        ]:
+            context[k] = data.get(k, None)
+
+    return context

tests/cli/resources/test_context.toml

@@ -0,0 +1,7 @@
+# Test TOML file for Hamilton CLI context loading
+
+# Define Hamilton headers as top-level values
+HAMILTON_CONFIG = {test_param = "test_value"}
+HAMILTON_FINAL_VARS = ["final_var1", "final_var2"]
+HAMILTON_INPUTS = {input_value = 42}
+HAMILTON_OVERRIDES = {override_value = "override"}

tests/cli/resources/test_tool_hamilton.toml

@@ -0,0 +1,8 @@
+# Test TOML file for Hamilton CLI context loading using [tool.hamilton] section
+
+[tool.hamilton]
+# Hamilton-specific configuration
+config = {test_param = "test_value"}
+final_vars = ["final_var1", "final_var2"]
+inputs = {input_value = 42, string_input = "test_string"}
+overrides = {override_value = "override"}

tests/cli/test_logic.py

@@ -93,3 +93,44 @@ def test_diff_node_versions():
     assert diff["reference_only"] == ["orders_per_customer"]
     assert diff["current_only"] == ["orders_per_distributor"]
     assert diff["edit"] == ["average_order_by_customer", "customer_summary_table"]
+
+
+def test_load_context_from_toml():
+    """Test loading context from a TOML file with top-level Hamilton headers."""
+    import os
+    os.environ["HAMILTON_CONFIG"] = "HAMILTON_CONFIG"
+    os.environ["HAMILTON_FINAL_VARS"] = "HAMILTON_FINAL_VARS"
+    os.environ["HAMILTON_INPUTS"] = "HAMILTON_INPUTS"
+    os.environ["HAMILTON_OVERRIDES"] = "HAMILTON_OVERRIDES"
+
+    toml_path = Path(__file__).parent / "resources" / "test_context.toml"
+
+    # Load context from TOML file
+    context = logic.load_context(toml_path)
+
+    # Check that the expected values are loaded
+    assert context["HAMILTON_CONFIG"] == {"test_param": "test_value"}
+    assert context["HAMILTON_INPUTS"] == {"input_value": 42}
+    assert context["HAMILTON_OVERRIDES"] == {"override_value": "override"}
+    # The TOML file has an array of final variables
+    assert context["HAMILTON_FINAL_VARS"] == ["final_var1", "final_var2"]
+
+
+def test_load_context_from_toml_tool_hamilton():
+    """Test loading context from a TOML file with [tool.hamilton] section."""
+    import os
+    os.environ["HAMILTON_CONFIG"] = "HAMILTON_CONFIG"
+    os.environ["HAMILTON_FINAL_VARS"] = "HAMILTON_FINAL_VARS"
+    os.environ["HAMILTON_INPUTS"] = "HAMILTON_INPUTS"
+    os.environ["HAMILTON_OVERRIDES"] = "HAMILTON_OVERRIDES"
+
+    toml_path = Path(__file__).parent / "resources" / "test_tool_hamilton.toml"
+
+    # Load context from TOML file with tool.hamilton section
+    context = logic.load_context(toml_path)
+
+    # Check that the expected values from [tool.hamilton] section are loaded
+    assert context["HAMILTON_CONFIG"] == {"test_param": "test_value"}
+    assert context["HAMILTON_INPUTS"] == {"input_value": 42, "string_input": "test_string"}
+    assert context["HAMILTON_OVERRIDES"] == {"override_value": "override"}
+    assert context["HAMILTON_FINAL_VARS"] == ["final_var1", "final_var2"]