From faeee504d006666370f53f620d1277879e159d9f Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:36:27 +0000
Subject: [PATCH 01/20] feat: Add Run object for manual hyperparameters and
 metrics tracking

---
 mthd/decorator.py | 76 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 59 insertions(+), 17 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 7180e50..755b67b 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -1,7 +1,7 @@
 import os
-
+from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, Optional, cast
+from typing import Any, Callable, Optional, cast
 
 from pydantic import BaseModel
 from rich.console import Console
@@ -10,6 +10,29 @@
 from mthd.domain.experiment import ExperimentRun
 from mthd.domain.git import StageStrategy
 from mthd.error import MthdError
+
+
+@dataclass
+class Run:
+    """Tracks experiment hyperparameters and metrics."""
+    _hypers: dict = None
+    _metrics: dict = None
+    
+    def set_hyperparameters(self, **kwargs) -> None:
+        """Set hyperparameters manually."""
+        self._hypers = kwargs
+        
+    def set_metrics(self, **kwargs) -> None:
+        """Set metrics manually."""
+        self._metrics = kwargs
+        
+    @property
+    def hyperparameters(self) -> Optional[dict]:
+        return self._hypers
+        
+    @property
+    def metrics(self) -> Optional[dict]:
+        return self._metrics
 from mthd.service.git import GitService
 from mthd.util.di import DI
 
@@ -30,23 +53,32 @@ def decorator(func: Callable) -> Callable:
         @wraps(func)
         def wrapper(*args, **kwargs):
             di = DI()
-            # @todo: handle this better (eg. positional args)
-            hyperparameters = cast(BaseModel, kwargs.get(hypers, None))
-            if not hyperparameters:
-                raise MthdError("Hyperparameters must be provided in the function call.")
+            run = Run()
             git_service = di[GitService]
-            # codebase_service = di[CodebaseService]
-
-            # Generate commit message
-
+            
             # Run experiment
-            metrics = func(*args, **kwargs)
+            metrics = func(*args, run=run, **kwargs)
+            
+            # Get experiment data from either the Run object or function args/return
+            if run.hyperparameters is not None:
+                hyper_dict = run.hyperparameters
+            else:
+                hyperparameters = cast(BaseModel, kwargs.get(hypers, None))
+                if not hyperparameters:
+                    raise MthdError("Hyperparameters must be provided either via Run object or function arguments")
+                hyper_dict = hyperparameters.model_dump()
+                
+            if run.metrics is not None:
+                metric_dict = run.metrics
+            else:
+                if not isinstance(metrics, BaseModel):
+                    raise MthdError("Metrics must be provided either via Run object or as BaseModel return value")
+                metric_dict = metrics.model_dump()
 
             experiment = ExperimentRun(
                 experiment=func.__name__,
-                hyperparameters=hyperparameters.model_dump(),
-                metrics=metrics.model_dump(),
-                # annotations=codebase_service.get_all_annotations(),
+                hyperparameters=hyper_dict,
+                metrics=metric_dict,
             )
             message = experiment.as_commit_message(template=template)
 
@@ -82,9 +114,19 @@ class Metrics(BaseModel):
         b: float
         c: str
 
+    # Example using function arguments/return
     @commit(hypers="hypers", template="run {experiment} at {timestamp}")
-    def test(hypers: Hyperparameters):
-        print("\n<Experiment goes here>\n")
+    def test1(hypers: Hyperparameters, run: Run):
+        print("\n<Experiment 1 goes here>\n")
         return Metrics(a=1, b=2.0, c="3")
 
-    test(hypers=Hyperparameters(a=1, b=2.0, c="3"))
+    # Example using Run object
+    @commit
+    def test2(run: Run):
+        print("\n<Experiment 2 goes here>\n")
+        run.set_hyperparameters(a=1, b=2.0, c="3")
+        run.set_metrics(a=1, b=2.0, c="3")
+        return None
+
+    test1(hypers=Hyperparameters(a=1, b=2.0, c="3"))
+    test2()

From e916d1af5454f350aed8052179bfe09f3f227ffd Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:42:51 +0000
Subject: [PATCH 02/20] refactor: Enhance commit decorator with optional
 context and improved hyperparameter handling

---
 mthd/decorator.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 755b67b..90138bd 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -1,7 +1,8 @@
 import os
+
 from dataclasses import dataclass
 from functools import wraps
-from typing import Any, Callable, Optional, cast
+from typing import Callable, Optional, cast
 
 from pydantic import BaseModel
 from rich.console import Console
@@ -10,31 +11,32 @@
 from mthd.domain.experiment import ExperimentRun
 from mthd.domain.git import StageStrategy
 from mthd.error import MthdError
+from mthd.service.git import GitService
+from mthd.util.di import DI
 
 
 @dataclass
 class Run:
     """Tracks experiment hyperparameters and metrics."""
+
     _hypers: dict = None
     _metrics: dict = None
-    
+
     def set_hyperparameters(self, **kwargs) -> None:
         """Set hyperparameters manually."""
         self._hypers = kwargs
-        
+
     def set_metrics(self, **kwargs) -> None:
         """Set metrics manually."""
         self._metrics = kwargs
-        
+
     @property
     def hyperparameters(self) -> Optional[dict]:
         return self._hypers
-        
+
     @property
     def metrics(self) -> Optional[dict]:
         return self._metrics
-from mthd.service.git import GitService
-from mthd.util.di import DI
 
 
 def commit(
@@ -43,6 +45,7 @@ def commit(
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
+    use_context: bool = False,
 ) -> Callable:
     """Decorator to auto-commit experimental code with scientific metadata.
 
@@ -55,10 +58,10 @@ def wrapper(*args, **kwargs):
             di = DI()
             run = Run()
             git_service = di[GitService]
-            
+
             # Run experiment
             metrics = func(*args, run=run, **kwargs)
-            
+
             # Get experiment data from either the Run object or function args/return
             if run.hyperparameters is not None:
                 hyper_dict = run.hyperparameters
@@ -67,7 +70,7 @@ def wrapper(*args, **kwargs):
                 if not hyperparameters:
                     raise MthdError("Hyperparameters must be provided either via Run object or function arguments")
                 hyper_dict = hyperparameters.model_dump()
-                
+
             if run.metrics is not None:
                 metric_dict = run.metrics
             else:

From 0465b412963b79f05443a12c479c4ccf1f3c0432 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:42:53 +0000
Subject: [PATCH 03/20] refactor: Modify commit decorator to support
 context-based and direct parameter approaches

---
 mthd/decorator.py | 38 ++++++++++++++++++++------------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 90138bd..11e26ba 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -56,30 +56,33 @@ def decorator(func: Callable) -> Callable:
         @wraps(func)
         def wrapper(*args, **kwargs):
             di = DI()
-            run = Run()
             git_service = di[GitService]
 
-            # Run experiment
-            metrics = func(*args, run=run, **kwargs)
-
-            # Get experiment data from either the Run object or function args/return
-            if run.hyperparameters is not None:
+            if use_context:
+                run = Run()
+                metrics = func(*args, run=run, **kwargs)
+                
+                if run.hyperparameters is None:
+                    raise MthdError("When using context, hyperparameters must be set via the Run object")
+                if run.metrics is None:
+                    raise MthdError("When using context, metrics must be set via the Run object")
+                
                 hyper_dict = run.hyperparameters
+                metric_dict = run.metrics
             else:
+                metrics = func(*args, **kwargs)
+                
                 hyperparameters = cast(BaseModel, kwargs.get(hypers, None))
                 if not hyperparameters:
-                    raise MthdError("Hyperparameters must be provided either via Run object or function arguments")
-                hyper_dict = hyperparameters.model_dump()
-
-            if run.metrics is not None:
-                metric_dict = run.metrics
-            else:
+                    raise MthdError("When not using context, hyperparameters must be provided as function arguments")
                 if not isinstance(metrics, BaseModel):
-                    raise MthdError("Metrics must be provided either via Run object or as BaseModel return value")
+                    raise MthdError("When not using context, metrics must be returned as a BaseModel")
+                
+                hyper_dict = hyperparameters.model_dump()
                 metric_dict = metrics.model_dump()
 
             experiment = ExperimentRun(
-                experiment=func.__name__,
+                experiment=func.__name__, 
                 hyperparameters=hyper_dict,
                 metrics=metric_dict,
             )
@@ -119,17 +122,16 @@ class Metrics(BaseModel):
 
     # Example using function arguments/return
     @commit(hypers="hypers", template="run {experiment} at {timestamp}")
-    def test1(hypers: Hyperparameters, run: Run):
+    def test1(hypers: Hyperparameters):
         print("\n<Experiment 1 goes here>\n")
         return Metrics(a=1, b=2.0, c="3")
 
-    # Example using Run object
-    @commit
+    # Example using Run context object
+    @commit(use_context=True)
     def test2(run: Run):
         print("\n<Experiment 2 goes here>\n")
         run.set_hyperparameters(a=1, b=2.0, c="3")
         run.set_metrics(a=1, b=2.0, c="3")
-        return None
 
     test1(hypers=Hyperparameters(a=1, b=2.0, c="3"))
     test2()

From 411d26280d5d6c867388655874385df71d8a1a42 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:49:29 +0000
Subject: [PATCH 04/20] refactor: Rename Run to Context and update method
 signatures

---
 mthd/decorator.py | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 11e26ba..90ea2cf 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -16,19 +16,19 @@
 
 
 @dataclass
-class Run:
+class Context:
     """Tracks experiment hyperparameters and metrics."""
 
-    _hypers: dict = None
-    _metrics: dict = None
+    _hypers: Optional[dict] = None
+    _metrics: Optional[dict] = None
 
-    def set_hyperparameters(self, **kwargs) -> None:
+    def set_hyperparameters(self, hypers: dict) -> None:
         """Set hyperparameters manually."""
-        self._hypers = kwargs
+        self._hypers = hypers
 
-    def set_metrics(self, **kwargs) -> None:
+    def set_metrics(self, metrics: dict) -> None:
         """Set metrics manually."""
-        self._metrics = kwargs
+        self._metrics = metrics
 
     @property
     def hyperparameters(self) -> Optional[dict]:
@@ -59,30 +59,30 @@ def wrapper(*args, **kwargs):
             git_service = di[GitService]
 
             if use_context:
-                run = Run()
+                run = Context()
                 metrics = func(*args, run=run, **kwargs)
-                
+
                 if run.hyperparameters is None:
                     raise MthdError("When using context, hyperparameters must be set via the Run object")
                 if run.metrics is None:
                     raise MthdError("When using context, metrics must be set via the Run object")
-                
+
                 hyper_dict = run.hyperparameters
                 metric_dict = run.metrics
             else:
                 metrics = func(*args, **kwargs)
-                
+
                 hyperparameters = cast(BaseModel, kwargs.get(hypers, None))
                 if not hyperparameters:
                     raise MthdError("When not using context, hyperparameters must be provided as function arguments")
                 if not isinstance(metrics, BaseModel):
                     raise MthdError("When not using context, metrics must be returned as a BaseModel")
-                
+
                 hyper_dict = hyperparameters.model_dump()
                 metric_dict = metrics.model_dump()
 
             experiment = ExperimentRun(
-                experiment=func.__name__, 
+                experiment=func.__name__,
                 hyperparameters=hyper_dict,
                 metrics=metric_dict,
             )
@@ -128,10 +128,10 @@ def test1(hypers: Hyperparameters):
 
     # Example using Run context object
     @commit(use_context=True)
-    def test2(run: Run):
+    def test2(context: Context):
         print("\n<Experiment 2 goes here>\n")
-        run.set_hyperparameters(a=1, b=2.0, c="3")
-        run.set_metrics(a=1, b=2.0, c="3")
+        context.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
+        context.set_metrics({"a": 1, "b": 2.0, "c": "3"})
 
     test1(hypers=Hyperparameters(a=1, b=2.0, c="3"))
     test2()

From 8f2f0cb6f46add4b46836285545f8d5b4f75f2df Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:49:30 +0000
Subject: [PATCH 05/20] refactor: Update context injection in decorator to
 improve UX

---
 mthd/decorator.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 90ea2cf..8b490e8 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -59,16 +59,17 @@ def wrapper(*args, **kwargs):
             git_service = di[GitService]
 
             if use_context:
-                run = Context()
-                metrics = func(*args, run=run, **kwargs)
+                context = Context()
+                kwargs['context'] = context
+                metrics = func(*args, **kwargs)
 
-                if run.hyperparameters is None:
-                    raise MthdError("When using context, hyperparameters must be set via the Run object")
-                if run.metrics is None:
-                    raise MthdError("When using context, metrics must be set via the Run object")
+                if context.hyperparameters is None:
+                    raise MthdError("When using context, hyperparameters must be set via the Context object")
+                if context.metrics is None:
+                    raise MthdError("When using context, metrics must be set via the Context object")
 
-                hyper_dict = run.hyperparameters
-                metric_dict = run.metrics
+                hyper_dict = context.hyperparameters
+                metric_dict = context.metrics
             else:
                 metrics = func(*args, **kwargs)
 

From 7128a19b4d4d82ea1fadc6203002bf8a9f9afa78 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:51:49 +0000
Subject: [PATCH 06/20] refactor: Update context usage to run in commit
 decorator

---
 mthd/decorator.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 8b490e8..284d338 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -59,17 +59,16 @@ def wrapper(*args, **kwargs):
             git_service = di[GitService]
 
             if use_context:
-                context = Context()
-                kwargs['context'] = context
-                metrics = func(*args, **kwargs)
+                run = Context()
+                metrics = func(*args, run=run, **kwargs)
 
-                if context.hyperparameters is None:
-                    raise MthdError("When using context, hyperparameters must be set via the Context object")
-                if context.metrics is None:
-                    raise MthdError("When using context, metrics must be set via the Context object")
+                if run.hyperparameters is None:
+                    raise MthdError("When using context, hyperparameters must be set via the Run object")
+                if run.metrics is None:
+                    raise MthdError("When using context, metrics must be set via the Run object")
 
-                hyper_dict = context.hyperparameters
-                metric_dict = context.metrics
+                hyper_dict = run.hyperparameters
+                metric_dict = run.metrics
             else:
                 metrics = func(*args, **kwargs)
 
@@ -129,7 +128,7 @@ def test1(hypers: Hyperparameters):
 
     # Example using Run context object
     @commit(use_context=True)
-    def test2(context: Context):
+    def test2(context: Context | None = None):
         print("\n<Experiment 2 goes here>\n")
         context.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
         context.set_metrics({"a": 1, "b": 2.0, "c": "3"})

From 5e0c131a69d3b5c2badd0e58bf31d91b4e4a7a12 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:51:51 +0000
Subject: [PATCH 07/20] refactor: Improve type hints for commit decorator with
 ParamSpec

---
 mthd/decorator.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 284d338..36a5001 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, Optional, cast
+from typing import Callable, Optional, ParamSpec, TypeVar, cast
 
 from pydantic import BaseModel
 from rich.console import Console
@@ -39,14 +39,17 @@ def metrics(self) -> Optional[dict]:
         return self._metrics
 
 
+P = ParamSpec('P')
+R = TypeVar('R')
+
 def commit(
-    fn: Optional[Callable] = None,
+    fn: Optional[Callable[P, R]] = None,
     *,
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: bool = False,
-) -> Callable:
+) -> Callable[P, R]:
     """Decorator to auto-commit experimental code with scientific metadata.
 
     Can be used as @commit or @commit(message="Custom message")
@@ -128,7 +131,7 @@ def test1(hypers: Hyperparameters):
 
     # Example using Run context object
     @commit(use_context=True)
-    def test2(context: Context | None = None):
+    def test2(context: Context):
         print("\n<Experiment 2 goes here>\n")
         context.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
         context.set_metrics({"a": 1, "b": 2.0, "c": "3"})

From 539fb6b7a839239bdb85bd8342ad428c86bf1f42 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:52:55 +0000
Subject: [PATCH 08/20] refactor: Fix ParamSpec typing and decorator signature
 for commit function

---
 mthd/decorator.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 36a5001..2cc5d72 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -39,23 +39,24 @@ def metrics(self) -> Optional[dict]:
         return self._metrics
 
 
-P = ParamSpec('P')
-R = TypeVar('R')
+P = ParamSpec("P")
+R = TypeVar("R")
+T = TypeVar("T")
 
 def commit(
-    fn: Optional[Callable[P, R]] = None,
+    fn: Optional[Callable[..., R]] = None,
     *,
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: bool = False,
-) -> Callable[P, R]:
+) -> Callable[..., R]:
     """Decorator to auto-commit experimental code with scientific metadata.
 
     Can be used as @commit or @commit(message="Custom message")
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: Callable[..., R]) -> Callable[..., R]:
         @wraps(func)
         def wrapper(*args, **kwargs):
             di = DI()

From 7bf073d8e941805de684e43ac8f0d74672255ae5 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:54:03 +0000
Subject: [PATCH 09/20] refactor: Fix context injection and type hints in
 decorator

---
 mthd/decorator.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 2cc5d72..d805cc0 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -63,16 +63,19 @@ def wrapper(*args, **kwargs):
             git_service = di[GitService]
 
             if use_context:
-                run = Context()
-                metrics = func(*args, run=run, **kwargs)
+                context = Context()
+                # Only inject if context isn't already provided
+                if 'context' not in kwargs:
+                    kwargs['context'] = context
+                metrics = func(*args, **kwargs)
 
-                if run.hyperparameters is None:
-                    raise MthdError("When using context, hyperparameters must be set via the Run object")
-                if run.metrics is None:
-                    raise MthdError("When using context, metrics must be set via the Run object")
+                if context.hyperparameters is None:
+                    raise MthdError("When using context, hyperparameters must be set via the Context object")
+                if context.metrics is None:
+                    raise MthdError("When using context, metrics must be set via the Context object")
 
-                hyper_dict = run.hyperparameters
-                metric_dict = run.metrics
+                hyper_dict = context.hyperparameters
+                metric_dict = context.metrics
             else:
                 metrics = func(*args, **kwargs)
 
@@ -130,7 +133,7 @@ def test1(hypers: Hyperparameters):
         print("\n<Experiment 1 goes here>\n")
         return Metrics(a=1, b=2.0, c="3")
 
-    # Example using Run context object
+    # Example using Context object
     @commit(use_context=True)
     def test2(context: Context):
         print("\n<Experiment 2 goes here>\n")
@@ -138,4 +141,4 @@ def test2(context: Context):
         context.set_metrics({"a": 1, "b": 2.0, "c": "3"})
 
     test1(hypers=Hyperparameters(a=1, b=2.0, c="3"))
-    test2()
+    test2(context=Context())  # Explicitly pass context to satisfy type checker

From 4c45440cb2eff6f99c8f1f2d8707293a15b69dd7 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:55:56 +0000
Subject: [PATCH 10/20] feat: Add overloaded commit decorator with improved
 type hints and UX

---
 mthd/decorator.py | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index d805cc0..6005b85 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, Optional, ParamSpec, TypeVar, cast
+from typing import Callable, Literal, Optional, ParamSpec, TypeVar, cast, overload
 
 from pydantic import BaseModel
 from rich.console import Console
@@ -43,6 +43,31 @@ def metrics(self) -> Optional[dict]:
 R = TypeVar("R")
 T = TypeVar("T")
 
+@overload
+def commit(
+    fn: None = None,
+    *,
+    hypers: str = "hypers",
+    template: str = "run {experiment}",
+    strategy: StageStrategy = StageStrategy.ALL,
+    use_context: Literal[True],
+) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
+
+@overload
+def commit(
+    fn: None = None,
+    *,
+    hypers: str = "hypers",
+    template: str = "run {experiment}",
+    strategy: StageStrategy = StageStrategy.ALL,
+    use_context: Literal[False] = False,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+@overload
+def commit(
+    fn: Callable[P, R],
+) -> Callable[P, R]: ...
+
 def commit(
     fn: Optional[Callable[..., R]] = None,
     *,
@@ -141,4 +166,4 @@ def test2(context: Context):
         context.set_metrics({"a": 1, "b": 2.0, "c": "3"})
 
     test1(hypers=Hyperparameters(a=1, b=2.0, c="3"))
-    test2(context=Context())  # Explicitly pass context to satisfy type checker
+    test2()  # No need to pass context

From ecfc83b6f8126ac59dd3926fc6e82d906eac09e2 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:58:16 +0000
Subject: [PATCH 11/20] refactor: Improve type hints and formatting in commit
 decorator

---
 mthd/decorator.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 6005b85..2421f8a 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -43,6 +43,7 @@ def metrics(self) -> Optional[dict]:
 R = TypeVar("R")
 T = TypeVar("T")
 
+
 @overload
 def commit(
     fn: None = None,
@@ -53,6 +54,7 @@ def commit(
     use_context: Literal[True],
 ) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
 
+
 @overload
 def commit(
     fn: None = None,
@@ -63,11 +65,13 @@ def commit(
     use_context: Literal[False] = False,
 ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
 
+
 @overload
 def commit(
     fn: Callable[P, R],
 ) -> Callable[P, R]: ...
 
+
 def commit(
     fn: Optional[Callable[..., R]] = None,
     *,
@@ -90,8 +94,8 @@ def wrapper(*args, **kwargs):
             if use_context:
                 context = Context()
                 # Only inject if context isn't already provided
-                if 'context' not in kwargs:
-                    kwargs['context'] = context
+                if "context" not in kwargs:
+                    kwargs["context"] = context
                 metrics = func(*args, **kwargs)
 
                 if context.hyperparameters is None:

From b8e80d1ec132eae83d1ba91b3c37751d1377b34f Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Mon, 10 Feb 2025 20:58:19 +0000
Subject: [PATCH 12/20] fix: Correct type hints for commit decorator to resolve
 pyright error

---
 mthd/decorator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 2421f8a..e54ed9d 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, Literal, Optional, ParamSpec, TypeVar, cast, overload
+from typing import Callable, Literal, Optional, ParamSpec, TypeVar, Union, cast, overload
 
 from pydantic import BaseModel
 from rich.console import Console
@@ -79,7 +79,7 @@ def commit(
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: bool = False,
-) -> Callable[..., R]:
+) -> Union[Callable[[Callable[..., R]], Callable[..., R]], Callable[..., R]]:
     """Decorator to auto-commit experimental code with scientific metadata.
 
     Can be used as @commit or @commit(message="Custom message")

From 4252e4623c3d9cd97863cff2f1d234a86f9db5c8 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 21:48:47 +0000
Subject: [PATCH 13/20] feat: Add e2e test for context-based experiment
 tracking with metrics and hyperparameters

---
 tests/e2e/e2e_test.py | 100 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 100 insertions(+)

diff --git a/tests/e2e/e2e_test.py b/tests/e2e/e2e_test.py
index bf6a688..73fb8f2 100644
--- a/tests/e2e/e2e_test.py
+++ b/tests/e2e/e2e_test.py
@@ -86,6 +86,39 @@ def train_model(hypers: Hyperparameters) -> Metrics:
         f.write(content)
 
 
+def create_context_experiment_file(temp_dir: Path, iteration: int) -> None:
+    """Create or update the experiment file using context-based API."""
+    content = f"""
+from mthd import commit
+from mthd.decorator import Context
+
+@commit(use_context=True)
+def train_model(context: Context):
+    # Set hyperparameters
+    context.set_hyperparameters({{
+        "learning_rate": 0.001 * ({iteration} + 1),
+        "batch_size": 32 * ({iteration} + 1),
+        "epochs": 10 * ({iteration} + 1)
+    }})
+    
+    # Simulate training
+    accuracy = 0.75 + ({iteration} * 0.05)  # Gradually improve accuracy
+    loss = 0.5 - ({iteration} * 0.1)        # Gradually decrease loss
+    
+    # Set metrics
+    context.set_metrics({{
+        "accuracy": accuracy,
+        "loss": max(0.1, loss)
+    }})
+
+if __name__ == "__main__":
+    train_model()
+"""
+
+    with open(temp_dir / "experiment.py", "w") as f:
+        f.write(content)
+
+
 def test_multiple_experiments(temp_dir: Path):
     """Test running multiple experiments and creating commits."""
     # Get the path to the Python executable in the virtual environment
@@ -151,3 +184,70 @@ def test_multiple_experiments(temp_dir: Path):
     output_lines = result.stdout.strip().split("\n")
     assert len(output_lines) > 0
     assert "Found 3 commit(s)" in output_lines[0]
+
+
+def test_multiple_context_experiments(temp_dir: Path):
+    """Test running multiple experiments using context-based API."""
+    # Get the path to the Python executable in the virtual environment
+    if os.name == "nt":  # Windows
+        python_path = temp_dir / ".venv" / "Scripts" / "python.exe"
+        mthd_path = temp_dir / ".venv" / "Scripts" / "mthd.exe"
+    else:  # Unix-like
+        python_path = temp_dir / ".venv" / "bin" / "python"
+        mthd_path = temp_dir / ".venv" / "bin" / "mthd"
+
+    # Run multiple iterations of the experiment
+    for i in range(4):
+        create_context_experiment_file(temp_dir, i)
+
+        # Run the experiment using the virtualenv python
+        subprocess.run(
+            [str(python_path), str(temp_dir / "experiment.py")],
+            check=True,
+        )
+
+    # Verify the commits
+    repo = git.Repo(temp_dir)
+    commits = list(repo.iter_commits())
+
+    # Should have 4 commits
+    assert len(commits) == 4
+
+    # All commits should be experiment commits
+    for commit in commits:
+        message = commit.message if isinstance(commit.message, str) else commit.message.decode("utf-8")
+
+        assert message.startswith("exp: ")
+        assert "metrics" in message
+        assert "hyperparameters" in message
+
+    # Test querying for experiments with high accuracy
+    result = subprocess.run(
+        [str(mthd_path), "query", "metrics.accuracy > 0.8"],
+        capture_output=True,
+        text=True,
+        cwd=temp_dir,
+    )
+    print(result.stdout)
+
+    assert result.returncode == 0
+
+    # Should find 2 commits (iterations 2 and 4 have accuracy > 0.8)
+    output_lines = result.stdout.strip().split("\n")
+    assert len(output_lines) > 0
+    assert "Found 2 commit(s)" in output_lines[0]
+
+    # Test querying for experiments with low loss
+    result = subprocess.run(
+        [str(mthd_path), "query", "metrics.loss < 0.5"],
+        capture_output=True,
+        text=True,
+        cwd=temp_dir,
+    )
+
+    assert result.returncode == 0
+
+    # Should find 3 commits (iterations 2, 3, and 4 have loss < 0.5)
+    output_lines = result.stdout.strip().split("\n")
+    assert len(output_lines) > 0
+    assert "Found 3 commit(s)" in output_lines[0]

From 94f8d431021dede0bcbca5a88bf48eaa99aad300 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 21:59:52 +0000
Subject: [PATCH 14/20] refactor: Update test2 function signature to include
 foo parameter

---
 mthd/decorator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index e54ed9d..4a7fc36 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -164,7 +164,7 @@ def test1(hypers: Hyperparameters):
 
     # Example using Context object
     @commit(use_context=True)
-    def test2(context: Context):
+    def test2(foo: int, context: Context):
         print("\n<Experiment 2 goes here>\n")
         context.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
         context.set_metrics({"a": 1, "b": 2.0, "c": "3"})

From 4a08032d05ec02f6203ea03bc42c25d31f245e76 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 21:59:54 +0000
Subject: [PATCH 15/20] feat: Improve type hints for commit decorator with
 Concatenate

---
 mthd/decorator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 4a7fc36..9878381 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import wraps
-from typing import Callable, Literal, Optional, ParamSpec, TypeVar, Union, cast, overload
+from typing import Callable, Concatenate, Literal, Optional, ParamSpec, TypeVar, Union, cast, overload
 
 from pydantic import BaseModel
 from rich.console import Console
@@ -52,7 +52,7 @@ def commit(
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: Literal[True],
-) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
+) -> Callable[[Callable[Concatenate[P, Context], R]], Callable[P, R]]: ...
 
 
 @overload

From 5fa9725068136a9be1442344d7cbcf5757dd0676 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 22:01:39 +0000
Subject: [PATCH 16/20] fix: Simplify decorator typing and remove conflicting
 function signature

---
 mthd/decorator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 9878381..9105ee1 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -52,7 +52,7 @@ def commit(
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: Literal[True],
-) -> Callable[[Callable[Concatenate[P, Context], R]], Callable[P, R]]: ...
+) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
 
 
 @overload
@@ -164,7 +164,7 @@ def test1(hypers: Hyperparameters):
 
     # Example using Context object
     @commit(use_context=True)
-    def test2(foo: int, context: Context):
+    def test2(context: Context):
         print("\n<Experiment 2 goes here>\n")
         context.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
         context.set_metrics({"a": 1, "b": 2.0, "c": "3"})

From c0eae371df90004282c9aaf7fb155a7da751197f Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 22:02:32 +0000
Subject: [PATCH 17/20] refactor: Improve type hints for commit decorator using
 ParamSpec and Concatenate

---
 mthd/decorator.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 9105ee1..4b5a01b 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -41,7 +41,6 @@ def metrics(self) -> Optional[dict]:
 
 P = ParamSpec("P")
 R = TypeVar("R")
-T = TypeVar("T")
 
 
 @overload
@@ -52,7 +51,7 @@ def commit(
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: Literal[True],
-) -> Callable[[Callable[..., R]], Callable[..., R]]: ...
+) -> Callable[[Callable[Concatenate[P, [Context]], R]], Callable[P, R]]: ...
 
 
 @overload
@@ -73,13 +72,13 @@ def commit(
 
 
 def commit(
-    fn: Optional[Callable[..., R]] = None,
+    fn: Optional[Callable[Concatenate[P, [Context]], R]] = None,
     *,
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
     use_context: bool = False,
-) -> Union[Callable[[Callable[..., R]], Callable[..., R]], Callable[..., R]]:
+) -> Union[Callable[[Callable[Concatenate[P, [Context]], R]], Callable[P, R]], Callable[P, R]]:
     """Decorator to auto-commit experimental code with scientific metadata.
 
     Can be used as @commit or @commit(message="Custom message")

From 77ce27b7f1782625be80449cd69aca5c4d8620b4 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 23:43:26 +0000
Subject: [PATCH 18/20] refactor: rename Context class to Run and improve
 decorator API semantics

---
 mthd/decorator.py     | 37 ++++++++++++++++++-------------------
 tests/e2e/e2e_test.py | 24 +++++++++++-------------
 2 files changed, 29 insertions(+), 32 deletions(-)

diff --git a/mthd/decorator.py b/mthd/decorator.py
index 4b5a01b..ba1fd1f 100644
--- a/mthd/decorator.py
+++ b/mthd/decorator.py
@@ -16,7 +16,7 @@
 
 
 @dataclass
-class Context:
+class Run:
     """Tracks experiment hyperparameters and metrics."""
 
     _hypers: Optional[dict] = None
@@ -41,6 +41,7 @@ def metrics(self) -> Optional[dict]:
 
 P = ParamSpec("P")
 R = TypeVar("R")
+T = TypeVar("T")
 
 
 @overload
@@ -50,8 +51,8 @@ def commit(
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
-    use_context: Literal[True],
-) -> Callable[[Callable[Concatenate[P, [Context]], R]], Callable[P, R]]: ...
+    implicit: Literal[False] = False,
+) -> Callable[[Callable[Concatenate[Run, P], R]], Callable[P, R]]: ...
 
 
 @overload
@@ -61,24 +62,24 @@ def commit(
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
-    use_context: Literal[False] = False,
+    implicit: Literal[True],
 ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
 
 
 @overload
 def commit(
     fn: Callable[P, R],
-) -> Callable[P, R]: ...
+) -> Callable[Concatenate[Run, P], R]: ...
 
 
 def commit(
-    fn: Optional[Callable[Concatenate[P, [Context]], R]] = None,
+    fn: Optional[Callable[..., R]] = None,
     *,
     hypers: str = "hypers",
     template: str = "run {experiment}",
     strategy: StageStrategy = StageStrategy.ALL,
-    use_context: bool = False,
-) -> Union[Callable[[Callable[Concatenate[P, [Context]], R]], Callable[P, R]], Callable[P, R]]:
+    implicit: bool = False,
+) -> Union[Callable[[Callable[..., R]], Callable[..., R]], Callable[..., R]]:
     """Decorator to auto-commit experimental code with scientific metadata.
 
     Can be used as @commit or @commit(message="Custom message")
@@ -90,11 +91,9 @@ def wrapper(*args, **kwargs):
             di = DI()
             git_service = di[GitService]
 
-            if use_context:
-                context = Context()
-                # Only inject if context isn't already provided
-                if "context" not in kwargs:
-                    kwargs["context"] = context
+            if not implicit:
+                context = Run()
+                args = [context] + list(args)
                 metrics = func(*args, **kwargs)
 
                 if context.hyperparameters is None:
@@ -156,17 +155,17 @@ class Metrics(BaseModel):
         c: str
 
     # Example using function arguments/return
-    @commit(hypers="hypers", template="run {experiment} at {timestamp}")
+    @commit(hypers="hypers", implicit=True, template="run {experiment} at {timestamp}")
     def test1(hypers: Hyperparameters):
         print("\n<Experiment 1 goes here>\n")
         return Metrics(a=1, b=2.0, c="3")
 
     # Example using Context object
-    @commit(use_context=True)
-    def test2(context: Context):
+    @commit()
+    def test2(run: Run, foo: int):
         print("\n<Experiment 2 goes here>\n")
-        context.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
-        context.set_metrics({"a": 1, "b": 2.0, "c": "3"})
+        run.set_hyperparameters({"a": 1, "b": 2.0, "c": "3"})
+        run.set_metrics({"a": 1, "b": 2.0, "c": "3"})
 
     test1(hypers=Hyperparameters(a=1, b=2.0, c="3"))
-    test2()  # No need to pass context
+    test2(5)  # No need to pass context
diff --git a/tests/e2e/e2e_test.py b/tests/e2e/e2e_test.py
index 73fb8f2..9cd58bc 100644
--- a/tests/e2e/e2e_test.py
+++ b/tests/e2e/e2e_test.py
@@ -62,7 +62,7 @@ class Metrics(BaseModel):
     accuracy: float
     loss: float
 
-@commit(hypers="hypers")
+@commit(hypers="hypers", implicit=True)
 def train_model(hypers: Hyperparameters) -> Metrics:
     # Simulate training
     accuracy = 0.75 + ({iteration} * 0.05)  # Gradually improve accuracy
@@ -86,16 +86,16 @@ def train_model(hypers: Hyperparameters) -> Metrics:
         f.write(content)
 
 
-def create_context_experiment_file(temp_dir: Path, iteration: int) -> None:
-    """Create or update the experiment file using context-based API."""
+def create_run_experiment_file(temp_dir: Path, iteration: int) -> None:
+    """Create or update the experiment file using run-based API."""
     content = f"""
 from mthd import commit
-from mthd.decorator import Context
+from mthd.decorator import Run
 
-@commit(use_context=True)
-def train_model(context: Context):
+@commit
+def train_model(run: Run):
     # Set hyperparameters
-    context.set_hyperparameters({{
+    run.set_hyperparameters({{
         "learning_rate": 0.001 * ({iteration} + 1),
         "batch_size": 32 * ({iteration} + 1),
         "epochs": 10 * ({iteration} + 1)
@@ -106,7 +106,7 @@ def train_model(context: Context):
     loss = 0.5 - ({iteration} * 0.1)        # Gradually decrease loss
     
     # Set metrics
-    context.set_metrics({{
+    run.set_metrics({{
         "accuracy": accuracy,
         "loss": max(0.1, loss)
     }})
@@ -161,7 +161,6 @@ def test_multiple_experiments(temp_dir: Path):
         text=True,
         cwd=temp_dir,
     )
-    print(result.stdout)
 
     assert result.returncode == 0
 
@@ -186,8 +185,8 @@ def test_multiple_experiments(temp_dir: Path):
     assert "Found 3 commit(s)" in output_lines[0]
 
 
-def test_multiple_context_experiments(temp_dir: Path):
-    """Test running multiple experiments using context-based API."""
+def test_multiple_run_experiments(temp_dir: Path):
+    """Test running multiple experiments using run-based API."""
     # Get the path to the Python executable in the virtual environment
     if os.name == "nt":  # Windows
         python_path = temp_dir / ".venv" / "Scripts" / "python.exe"
@@ -198,7 +197,7 @@ def test_multiple_context_experiments(temp_dir: Path):
 
     # Run multiple iterations of the experiment
     for i in range(4):
-        create_context_experiment_file(temp_dir, i)
+        create_run_experiment_file(temp_dir, i)
 
         # Run the experiment using the virtualenv python
         subprocess.run(
@@ -228,7 +227,6 @@ def test_multiple_context_experiments(temp_dir: Path):
         text=True,
         cwd=temp_dir,
     )
-    print(result.stdout)
 
     assert result.returncode == 0
 

From 26d60025c2055c74752392a54735064ae42b401f Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 23:48:58 +0000
Subject: [PATCH 19/20] docs: update README with new Run API and improve code
 example clarity

---
 README.md | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index d09b338..14c20e5 100644
--- a/README.md
+++ b/README.md
@@ -25,7 +25,7 @@ Every time you run an experiment, your code will be auto-committed with metadata
 * Query your scientific log, e.g. `mthd query metrics.accuracy < 0.8`.
 
 ```python
-from mthd import commit
+from mthd import commit, Run
 from pydantic import BaseModel
 
 class Hypers(BaseModel):
@@ -39,15 +39,18 @@ class Metrics(BaseModel):
     accuracy: float
 
 
-@commit(hypers="hypers")
-def my_experiment(hypers: Hypers) -> Metrics:
+@commit
+def my_experiment(run: Run) -> Metrics:
     ...
-    # experiment
+
+    run.set_hyperparameters({ ... })
+
     ...
 
-    metrics = Metrics(...)
+    run.set_metrics({ ... })
 
-    return metrics
+if __name__ == "__main__":
+    my_experiment(Hypers(lr=0.001, epochs=100))
 ```
 
 Then run your experiment:

From 4bbc56fe379ddcde3edb768b87180c8ca4fb5093 Mon Sep 17 00:00:00 2001
From: Rory Byrne <rory@rory.bio>
Date: Fri, 14 Feb 2025 23:52:06 +0000
Subject: [PATCH 20/20] docs: simplify example code in README by removing
 unnecessary parameters

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 14c20e5..4da627d 100644
--- a/README.md
+++ b/README.md
@@ -50,7 +50,7 @@ def my_experiment(run: Run) -> Metrics:
     run.set_metrics({ ... })
 
 if __name__ == "__main__":
-    my_experiment(Hypers(lr=0.001, epochs=100))
+    my_experiment()
 ```
 
 Then run your experiment: