refactor: Replace ad-hoc dispatch with custom @singledispatch

[dangotbanned]

Description

Adds a slimmed-down version of @functools.singledispatch, with the initial use replacing _supertyping._same_supertype in #3396.

I can see a few other places (particularly DType-related) that could later benefit.
So, I went ahead and added more tests + docs than I would normally do for an internal tool.

Why not @functools.singledispatch?

Most of the stdlib code is dedicated to two features I don't want to use.

mro-based dispatch

That allows you to register ABCs (like Mapping, Iterable, etc) or any regular class and have it's subclasses match.

Definitely pretty clever stuff, but I'd prefer to have a simpler version that needs every class to be added explicitly

• https://github.com/python/cpython/blob/cf71e34940e2314ee7ca00961d86a7172286eeea/Lib/functools.py#L756-L894
• https://github.com/python/cpython/blob/cf71e34940e2314ee7ca00961d86a7172286eeea/Lib/functools.py#L921-L926

Registration via annotation

Forward references are resolved eagerly¹ and only classes and unions of classes are supported.

IMO, that's just a complicated way of accepting *types: type[Any] in the decorator - so I did that instead 😅

• https://github.com/python/cpython/blob/cf71e34940e2314ee7ca00961d86a7172286eeea/Lib/functools.py#L937-L941
• https://github.com/python/cpython/blob/cf71e34940e2314ee7ca00961d86a7172286eeea/Lib/functools.py#L959-L991

Related issues

• Child of feat(DRAFT): Add get_supertype #3396

Footnotes

1. ibis has @lazy_singledispatch that appears to solve this issue ↩

[dangotbanned]

I've left out validating the types for now.

• 3.9: @functools.singledispatch supports type[Any]
• 3.10: Same support, but that's the version which introduces types.UnionType
• 3.11: @functools.singledispatch supports typing.Union | types.UnionType | type[Any]

Since #3204 is not resolved, I'd rather leave this runtime check until then

[dangotbanned]

@FBruzzesi
I'm not sure if (#3396 (comment)) got lost, but I thought I'd give convincing you one last try 🙂

#3410 is a beautiful solution, but, at the same time, quite intense to maintain.
It requires it's dedicated test suite

I think this was a documentation issue on my part.
I tried out (literate-style tests), but as you can see from (revert: Keep only what is required) - removing all of that left only 2 lines uncovered.

Do you still have the opinion that (https://github.com/narwhals-dev/narwhals/blob/308f389f2e68079053e9c659394cdd0d4207410d/narwhals/_dispatch.py) is intense to maintain?

---

and it serves a single purpose in the codebase, which makes me feel a bit uneasy on it.

I thought replacing lots of other non-supertyping stuff wouldn't be a good idea in this PR, since it was based on dtypes/supertyping.

Here's three unrelated use-cases (dtypes/supertyping-dispatch...@just_dispatch-use-cases).
Only one was what I had planned in the PR description:

I can see a few other places (particularly DType-related) that could later benefit.

I do feel the need to stress that these are not the only places that can use it! 😅
Just didn't wanna go replacing too much if it is a hard no

[FBruzzesi]

Thanks @dangotbanned - I like this implementation much more than the first one.

I can see the benefits, and removing _SupertypeCase is already a big advantage I think.

It's a bit unclear how the typing issues get fixed with this approach, but I guess I will need an editor to better understand all the details 🙈

[FBruzzesi]

Thanks @dangotbanned - last comment to resolve is #3410 (comment)

---

One additional comment I have, that might be worth adding in just_dispatch docs is that we need to make sure that this is used only internally (i.e. functions that are not publicly exposed) as we don't do anything to maintain information such as docstrings. I didn't give it a try in an editor, but I would expect that to be the case as we go from a function to an instance wrapping such function.

[FBruzzesi]

Just to double check with you: there is no type checker error because regardless of the input argument annotation (DType), after being decorated, all we care about is the return type in the annotation (i.e. JustDispatch[str]).

[dangotbanned]

@FBruzzesi

Just to double check with you: there is no type checker error because regardless of the input argument annotation (DType), after being decorated,

Yeah that is to be expected (at least partially)

So the @just_dispatch-decorated function gets this signature:

narwhals/narwhals/_dispatch.py

Lines 50 to 52 in ddd5d98

	def __call__(self, arg: object, *args: Any, **kwds: Any) -> R_co:
	"""Dispatch on the type of the first argument, passing through all arguments."""
	return self.dispatch(arg.__class__)(arg, *args, **kwds)

I actually made it stricter than [@functools.singledispatch], which doesn't require¹ an argument:

    def __call__(self, /, *args: Any, **kwargs: Any) -> _T: ...

Besides that tweak, I need to note this only applies to the initally decorated function.
You can see in the old tests how all others still work as expected

Why?

I had initially planned to use upper_bound to derive the type, but removed it in (02b7811).

Partially that was to get mypy off my back 😅 - however the more I thought about it - the only accurate type is object.

Here's a contrived example, but it shows that you can register any type:

Show code

import polars as pl
from polars.datatypes import classes as pl_dtypes

import narwhals as nw
from narwhals._dispatch import just_dispatch
from narwhals._utils import qualified_type_name
from narwhals.dtypes import DType
from narwhals.dtypes._utils import validate_into_dtype
from narwhals.typing import IntoDType  # noqa: TC001


@just_dispatch
def into_dtype(obj: IntoDType) -> DType:
    if isinstance(obj, DType):
        return obj
    validate_into_dtype(obj)
    return obj()


@into_dtype.register(pl_dtypes.DataTypeClass)
def from_polars_meta(obj: type[pl.DataType] | pl_dtypes.DataTypeClass) -> DType:
    return into_dtype(getattr(nw, obj.__name__))


@into_dtype.register(
    *pl_dtypes.SignedIntegerType.__subclasses__(),
    *pl_dtypes.UnsignedIntegerType.__subclasses__(),
    pl.Float32,
    pl.Float64,
    pl.String,
    pl.Boolean,
    pl.Binary,
    pl.Date,
    pl.Time,
    pl.Categorical,
    pl.Unknown,
    pl.Object,
)
def from_polars(obj: pl_dtypes.DataType) -> DType:
    return from_polars_meta(obj.base_type())


ok_1 = into_dtype(nw.Date)
ok_2 = into_dtype(nw.Date())
ok_3 = into_dtype(pl.Date)
ok_4 = into_dtype(pl.Date())


print("\n".join(qualified_type_name(dtype) for dtype in [ok_1, ok_2, ok_3, ok_4]))

narwhals.dtypes._classes.Date
narwhals.dtypes._classes.Date
narwhals.dtypes._classes.Date
narwhals.dtypes._classes.Date

I didn't declare the function supported some types, but made it possible to do so later 😄

Footnotes

1. It does runtime validation to enforce >=1 argument, but that isn't helpful for typing 😢 ↩

[FBruzzesi]

Really nice test with Datetime distinction between main and v1 👌🏼