[Feature]: Fix position of StepContext in decorated class methods

[ngr]

What would you like?

The current implementation of the @durable_step decorator injects the StepContext object as the first positional argument, moving self (instanse pointer) to the second position if it decorates an method inside a class. This contradicts the PEP8 requirements:

Always use self for the first argument to instance methods.
Always use cls for the first argument to class methods.

I suggest refactoring the decorator behavior while it's not too late, although this will become a breaking change for the currently decorated class methods already running.

I am ready to try to solve this if approved.

Is this a breaking change?

Yes

[leandrodamascena]

I think the breaking change might be worth considering here. We could try to detect self/cls via inspect and adjust the argument order automatically, but honestly that feels a bit fragile to me - it depends on the parameter being literally named self or cls, and if someone uses a different name it just breaks silently.

For me this isn't just a design preference thing - the current behavior will actually break type checkers, potentially mess up IDE autocomplete, and might confuse developers who naturally expect self to be the first argument in methods.

I'm generally not a fan of changes that break compatibility, especially when it comes to the Semantic Versioning framework, as they always cause disruption for users. But perhaps we could make some kind of announcement or add a deprecation notice first to give people time to migrate?

Honestly, I don't know what the best approach would be here, and I'd like to hear other opinions on the matter.

cc @yaythomas

[yaythomas]

durable_step was originally conceived for functions, not methods.

if we want to make it work with instance methods, rather than brittle inspection on self/cls, we could do something like this, which maintains backwards compatibility with the current decorator.

(WARNING: I've not particularly tested this extensively outside of a quick smoke test and I've not fully thought it through yet - please feel free assume there are errors here and correct me if I overlooked something)

class durable_step:
    """Decorator that wraps a callable for use in durable step operations.
    
    This decorator uses the descriptor protocol to automatically detect whether
    it's decorating a plain function or a class method, ensuring correct parameter
    ordering in both cases.
    
    For plain functions:
        @durable_step
        def my_step(step_context: StepContext, arg: str) -> str:
            return f"Result: {arg}"
        
        # Usage: context.step("my_step", my_step(arg="value"))
    
    For class methods:
        class MyClass:
            @durable_step
            def my_method(self, step_context: StepContext, arg: str) -> str:
                return f"{self.value}: {arg}"
        
        # Usage: obj = MyClass()
        #        context.step("my_method", obj.my_method(arg="value"))
    
    The decorator ensures that:
    - For functions: StepContext is injected as the first parameter
    - For methods: self is preserved as the first parameter, StepContext as second
    - The _original_name attribute is set for operation identification
    """
    
    def __init__(self, func: Callable[Concatenate[StepContext, Params], T]):
        """Initialize the descriptor with the wrapped function.
        
        Args:
            func: The function to wrap for durable step execution
        """
        self.func = func
    
    def __get__(
        self, 
        obj: object | None, 
        objtype: type | None = None
    ) -> Callable[Params, Callable[[StepContext], T]]:
        """Called when the decorator is accessed as a method attribute.
        
        This method is invoked by Python's descriptor protocol when the decorated
        function is accessed as an attribute of a class instance (e.g., obj.method).
        
        Args:
            obj: The instance the method is being accessed from (None for class access)
            objtype: The class type the method belongs to
            
        Returns:
            A wrapper function that correctly handles method binding
        """
        if obj is None:
            # Accessed on class, not instance - return self for further binding
            return self  # type: ignore
        
        # Method access - obj is the instance (self)
        # Return a wrapper that preserves self in the correct position
        def wrapper(*args: Params.args, **kwargs: Params.kwargs) -> Callable[[StepContext], T]:
            def function_with_arguments(context: StepContext) -> T:
                # Call with correct order: self, context, *args, **kwargs
                return self.func(obj, context, *args, **kwargs)  # type: ignore
            
            function_with_arguments._original_name = self.func.__name__  # type: ignore
            return function_with_arguments
        
        return wrapper
    
    def __call__(
        self, 
        *args: Params.args, 
        **kwargs: Params.kwargs
    ) -> Callable[[StepContext], T]:
        """Called when the decorator is used on a plain function.
        
        This method is invoked when the decorated function is called directly
        (not as a method attribute), which happens for plain functions.
        
        Args:
            *args: Positional arguments to bind to the function
            **kwargs: Keyword arguments to bind to the function
            
        Returns:
            A function that takes StepContext and executes the wrapped function
        """
        def function_with_arguments(context: StepContext) -> T:
            # Call with correct order: context, *args, **kwargs
            return self.func(context, *args, **kwargs)
        
        function_with_arguments._original_name = self.func.__name__  # type: ignore
        return function_with_arguments

[yaythomas]

above won't work with @classmethod, but I'm tempted NOT to include this - I can't quickly think of a solid way of implementing this outside of making (brittle) assumptions about the first parameter (if args and isinstance(args[0], type)).

i.e only support functions and instance methods.

Can also consider distinct/separate decorator for this if the need exists (i.e @durable_step_classmethod vs @durable_step.

[ngr]

I like the implementation by @yaythomas
And the idea of giving up for now the implementation of support for @classmethod with a separate decorator is also okay.

Looks like this is also applicable for all the three: durable_with_child_context and durable_wait_for_callback.

[yaythomas]

yes, if we do this it should be for all the decorators.

The reason we made it function only (not method), is because using a class very naturally will lead a developer to maintain state on the class - in fact, you could argue in Python one of the prime motivations for using a class + methods over a function is because you want to maintain state between method calls.

I worry that this will very easily lead consumers down a path of writing non-deterministic code. Of course, you could argue that we already have the firm principle that all non-deterministic code must be inside steps and that state should be stored as the result of a step operation to retain it across re-plays, but even so, having this decorator work on instance methods makes it easier (and even makes it natural) to do the wrong thing.

To illustrate the point with a simple example:

class MyClass:
    def __init__(self):
        self.counter = 1

    @durable_step
    def one(self):
        self.counter = self.counter + 1

    @durable_step
    def two(self):
        self.counter = self.counter + 1


@durable_execution
def my_handler(event, context: DurableContext) -> str:
    obj = MyClass()

    # this adds 1 to the counter only on the FIRST execution
    context.step(obj.one)

    # function suspends
    context.wait(Duration.from_minutes(5))

    # When function wakes up again after the wait, "one" will not execute again, because it's already completed
    context.step(obj.two)

    # returns 2, not 3. this is surprising.
    return obj.counter

@ngr do you have particular use cases in mind for using a class? While we're considering this enhancement, the purpose for why a class might be useful could help guide us to a solution :-)

Regardless of what we land on here, we should do a better job documenting the decorator at present only supports functions and the implications of not storing state in the context for replays.

[ngr]

@yaythomas I have a use-case indeed. :)

Here is a very simple example where I want to keep the instance of the class not only through steps, but through different invocations of the Lambda:
https://github.com/ngr/aws-durable-execution-sdk-python/blob/sandbox/sandbox/src/app.py

This approach was implemented in the sosw framework, about 8 years ago and is still quite useful.
I understand that proper documentation and guidelines are critical, so if you approve I would be happy to clean up and open a PR.

P.S. I don't think that allowing the decorator to support instance methods would encourage people to write non-deterministic code. If somebody has already figured out how to store artifacts in the Lambda sandbox (container) cross-invocation, he has already learned all the mistakes there.