Extending subscriptable function allowing overloaded/polymorphic functions

I’d like to propose overloaded/polymorphic functions based on using PEP 718 - Subscriptable functions to resolve which Callable is used during type checking and runtime.

Let’s take a look at this use case where I have an object that is immutable but supports mutability by modifying private state.

class ImmutableObject:
    def __init__(self, x):
        self._mutable = True
        self.x = x
        self._mutable = False

    def __setattr__(self, __name: str, __value: Any) -> None:
        if __name.startswith("_") or self._mutable:
            super().__setattr__(__name, __value)
        else:
            raise TypeError(f"Cannot set attribute {__name} on immutable object")

    def __eq__(self, other: object) -> bool:
        return isinstance(other, ImmutableObject) and self.x == other.x

I define two copy functions, one classic and one that allows mutation of data within a context block:

def _copy(obj: ImmutableObject) -> ImmutableObject:
    copied = ImmutableObject(obj.x)
    return copied

@contextmanager
def _copy_in_mutable_context(obj: ImmutableObject) -> Iterator[ImmutableObject]:
    copied = ImmutableObject(obj.x)
    try:
        copied._mutable = True
        yield copied
    finally:
        copied._mutable = False

The goal is then to define a copy() function that can choose between the two functions based on the subscriptable function syntax.

Proposed syntax

def copy[R](obj: ImmutableObject) -> Callable[..., R]:
    function_getter = {
        Iterator: _copy_in_mutable_context,
        ImmutableObject: _copy
    }
    return function_getter[R](obj)   # <-- R is available at runtime

foobar= ImmutableObject(1)
with copy[Iterator](foobar) as copied:
    copied.x = 2
    assert copied == ImmutableObject(2)

Implementation with Python <= 3.12

Expand

Currently this is possible using a map:

copy = {
    Iterator: _copy_in_mutable_context,
    ImmutableObject: _copy
}

with copy[Iterator](foobar) as copied:
    ...

Or adding a flag/switch to select the right Callable:

def copy[R](obj: ImmutableObject, rtype: type[R]) -> Callable[..., R]:
    function_getter = {
        Iterator: _copy_in_mutable_context,
        ImmutableObject: _copy
    }
    return function_getter[rtype](obj)

with copy[Iterator](foobar, Iterator) as copied:
    ...

Similar use case:

with my_create(a=10) as new_obj:
    pass

new_obj.a  # => 10
new_obj.is_saved()  # => True

In the above example, I’ve passed an empty “block” to my my_create function. Things work as expected (my_obj was initialized, and saved), but the formatting looks a little wonky, and the with block seems unnecessary.