From 4a0694e24a5566c7f59130bc83e565f2d28c411d Mon Sep 17 00:00:00 2001
From: Jan Hurst <jan@janhurst.com>
Date: Tue, 23 Apr 2024 07:46:31 +0800
Subject: [PATCH] add a "resolve from config" wrapper around the resolve
 decorator (#828)

Add a small wrapper around resolve function modifier to "resolve from config".

This is a quality of life addition from Jan Hurst! Something he uses and helps
to cut down on the purposeful verbosity of the prior way.

* add resolve from config wrapper

* fix import lint

* add some docstrings
---
 hamilton/function_modifiers/__init__.py  |  1 +
 hamilton/function_modifiers/delayed.py   | 15 +++++++++++++++
 tests/function_modifiers/test_delayed.py | 21 ++++++++++++++++++++-
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/hamilton/function_modifiers/__init__.py b/hamilton/function_modifiers/__init__.py
index 2bfcc779b..1f52edcb2 100644
--- a/hamilton/function_modifiers/__init__.py
+++ b/hamilton/function_modifiers/__init__.py
@@ -87,6 +87,7 @@
 
 resolve = delayed.resolve
 ResolveAt = delayed.ResolveAt
+resolve_from_config = delayed.resolve_from_config
 
 # materialization stuff
 load_from = adapters.load_from
diff --git a/hamilton/function_modifiers/delayed.py b/hamilton/function_modifiers/delayed.py
index d2bd04b7d..c4b8b9659 100644
--- a/hamilton/function_modifiers/delayed.py
+++ b/hamilton/function_modifiers/delayed.py
@@ -148,3 +148,18 @@ def resolve(self, config: Dict[str, Any], fn: Callable) -> NodeTransformLifecycl
             if key in config:
                 kwargs[key] = config[key]
         return self.decorate_with(**kwargs)
+
+
+class resolve_from_config(resolve):
+    """Decorator class to delay evaluation of decorators until after the configuration is available.
+    Note: this is a power-user feature, and you have to enable power-user mode! To do so, you have
+    to add the configuration hamilton.enable_power_user_mode=True to the config you pass into the
+    driver.
+
+    This is a convenience decorator that is a subclass of `resolve` and passes
+    `ResolveAt.CONFIG_AVAILABLE` to the `when` argument such that the decorator is resoled at
+    compile time, E.G. when the driver is instantiated.
+    """
+
+    def __init__(self, *, decorate_with: Callable[..., NodeTransformLifecycle]):
+        super().__init__(when=ResolveAt.CONFIG_AVAILABLE, decorate_with=decorate_with)
diff --git a/tests/function_modifiers/test_delayed.py b/tests/function_modifiers/test_delayed.py
index 6940de39e..787dffd74 100644
--- a/tests/function_modifiers/test_delayed.py
+++ b/tests/function_modifiers/test_delayed.py
@@ -3,7 +3,13 @@
 import pytest
 
 from hamilton import settings
-from hamilton.function_modifiers import ResolveAt, base, extract_columns, resolve
+from hamilton.function_modifiers import (
+    ResolveAt,
+    base,
+    extract_columns,
+    resolve,
+    resolve_from_config,
+)
 
 CONFIG_WITH_POWER_MODE_ENABLED = {
     settings.ENABLE_POWER_USER_MODE: True,
@@ -58,6 +64,19 @@ def test_dynamic_resolves():
     assert decorator_resolved.columns == ("a", "b")
 
 
+def test_dynamic_resolve_with_configs():
+    decorator = resolve_from_config(
+        decorate_with=lambda cols_to_extract: extract_columns(*cols_to_extract),
+    )
+    decorator_resolved = decorator.resolve(
+        {"cols_to_extract": ["a", "b"], **CONFIG_WITH_POWER_MODE_ENABLED}, fn=test_dynamic_resolves
+    )
+    # This uses an internal component of extract_columns
+    # We may want to add a little more comprehensive testing
+    # But for now this will work
+    assert decorator_resolved.columns == ("a", "b")
+
+
 def test_dynamic_fails_without_power_mode_fails():
     decorator = resolve(
         when=ResolveAt.CONFIG_AVAILABLE,