Allow decorator use for functions defined in other modules

[skrawcz]

Is your feature request related to a problem? Please describe.
The following code throws errors:

Under: my_module/
__init__.py

from .base import a, z
from .base_extended import b_p, d, e

base.py

def a(input: int) -> int:
    return input * 2

def z(input: int) -> int:
    return input * 3

base_extended.py

from hamilton.function_modifiers import parameterize, value

from my_module import base

b_p = parameterize(
    b={"input":value(1)},
    c={"input":value(2)} 
)(base.a) 

def d(b: int, c: int) -> int:
    return b + c

Then in my run.py outside my_module.

from hamilton import driver
import my_module
dr = driver.Builder().with_modules(my_module).build()
dr.display_all_functions("base_extended.png")
print(dr.execute(['b', 'c']))

Throws two possible things (if you allow module overrides one error goes away):

1. b is redefined twice.
2. a is missing from the graph.

Describe the likely problem

1. I think we're modifying the underlying function in the parameterize decorator which results in a being lost.
2. Because of (2) we have a effectively as a function twice. Once via the parameterize decorator, once via module import. This then means that we're defining b and c twice.

Describe the solution you'd like

1. we should be able to import functions and use them in a decorator - without impacting the use of that function in hamilton as a standalone by itself thing.
2. with (1) resolved I think things should then just work.

Describe alternatives you've considered
N/A

Additional context
Comes from a problem a user ran into in slack.

[skrawcz]

Use cases to think through:

def _some_func(...):
   # ...

p = parameterize(...)(_some_func) # this doesn't work.
p2 = parameterize(... some other parameterization)(_some_func) # this doesn't work.

def some_func(...): # we want it in the graph
   # ...

p = parameterize(...)(some_func) # this doesn't work. - redefinition
p2 = parameterize(... some other parameterization)(some_func) # this doesn't work - only assume one parameterization of it

# some module
def some_func():...
   # ...

# another module
import some_module

p = parameterize(...)(some_module.some_func) # this modifies `some_func`  and `some_func` cannot be part of a graph explicitly.
p2 = parameterize(... some other parameterization)(some_module.some_func) # this doesn't work - only assume one parameterization of it

[skrawcz]

For example, I have a function that receives a list of torch.Dataset and generates a list of DataLoader, and I would like each project-module to be able to define its node using this function, since the nodes it depends on and the nodes that used this output will be different. The same with a function to generate plots from a list of arrays, which I would like to be able to create nodes with it for each dataflow.

[EltonCN]

A temporary workaround (the __call__ bellow summarizes what the class does):

import inspect
from functools import wraps
from typing import Any

from docstring_parser import compose, parse
from hamilton.function_modifiers import LiteralDependency, UpstreamDependency


class function_variation():
    def __init__(self, param_map, new_name: str | None = None):
        upstream_dependencies: dict[str, str] = {}
        literal_dependencies: dict[str, Any] = {}

        for param_name in param_map:
            target = param_map[param_name]

            if isinstance(target, UpstreamDependency):
                upstream_dependencies[param_name] = target.source
            elif isinstance(target, LiteralDependency):
                literal_dependencies[param_name] = target.value
            else:
                raise ValueError(
                    f"Parameter map must be a value or source. Parameter {param_name} is {type(target)}")

        self.upstream_dependencies = upstream_dependencies
        self.literal_dependencies = literal_dependencies

        self.new_name = new_name

    def _change_parameters(self, func):
        name = func.__name__
        if self.new_name is not None:
            name = self.new_name

        @wraps(func)
        def wrapper(*args, **kwargs):

            for literal_dep_name in self.literal_dependencies:
                if literal_dep_name in kwargs:
                    raise TypeError(
                        f"{name}() got an unexpected keyword argument '{literal_dep_name}'")

            for old_param, new_param in self.upstream_dependencies.items():
                if new_param in kwargs:
                    kwargs[old_param] = kwargs.pop(new_param)

            kwargs.update(self.literal_dependencies)

            return func(*args, **kwargs)

        return wrapper

    def _change_name(self, func):
        if self.new_name is not None:
            func.__name__ = self.new_name

            qualname = func.__qualname__.split(".")
            qualname[-1] = self.new_name

            func.__qualname__ = ".".join(qualname)

    def _update_signature(self, func):
        signature = inspect.signature(func)
        params = list(signature.parameters.values())

        new_params = []
        for param in params:
            if param.name in self.upstream_dependencies:
                new_name = self.upstream_dependencies[param.name]
                new_param = param.replace(name=new_name)

                new_params.append(new_param)
            elif param.name not in self.literal_dependencies:
                new_params.append(param)

        new_signature = signature.replace(parameters=new_params)
        func.__signature__ = new_signature

    def _update_docstring(self, func):
        docstring: str | None = func.__doc__

        if docstring is None:
            return

        # Rename parameters
        for old_name in self.upstream_dependencies:
            new_name = self.upstream_dependencies[old_name]

            docstring = docstring.replace(old_name, new_name)

        # Remove literal parameters
        docstring_obj = parse(docstring)
        remove_param = [
            param for param in docstring_obj.params if param.arg_name in self.literal_dependencies]
        docstring_obj.meta = [
            meta for meta in docstring_obj.meta if meta not in remove_param]
        docstring = compose(docstring_obj)

        func.__doc__ = docstring

    def _update_module(self, func):

        # 0 = _update_module|1 = __call__ NOSONAR
        func.__module__ = inspect.getmodule(inspect.stack()[2][0]).__name__

    def _update_annotations(self, func):

        func.__annotations__ = func.__annotations__.copy()

        for old_name in self.upstream_dependencies:
            new_name = self.upstream_dependencies[old_name]
            func.__annotations__[new_name] = func.__annotations__[old_name]
            del func.__annotations__[old_name]

        for name in self.literal_dependencies:
            del func.__annotations__[name]

    def __call__(self, func):
        func = self._change_parameters(func)
        self._change_name(func)
        self._update_signature(func)
        self._update_docstring(func)
        self._update_module(func)
        self._update_annotations(func)

        return func

and using as:

from hamilton.function_modifiers import parameterize, value

from my_module.base import a

b = function_variation({"input":value(1)}, "b")(a)
c = function_variation({"input":value(2)}, "c")(a)

def d(b: int, c: int) -> int:
    return b + c

Where the __init__.py becomes:

from .base import a, z
from .base_extended import b, c, d

I tested and it worked, and the generated variations/parameterized functions can also be used independently.

[skrawcz]

@EltonCN nice.

I'm thinking though there could be a simpler approach - #1282 - I haven't looked at all test failures but it seems to get pretty far...