feat(typing): Make Implementation less opaque

[dangotbanned]

What type of PR is this? (check all applicable)

• 💾 Refactor
• ✨ Feature
• 🐛 Bug Fix
• 🔧 Optimization
• 📝 Documentation
• ✅ Test
• 🐳 Other

Related issues

• feat(typing): Make IntoBackend generic #3002
• Based on (@overload interacts badly with @property even with the type: ignore workaround python/mypy#9937)
• Issue raised privately by @FBruzzesi
  • Explained in (feat(typing): Make Implementation less opaque #3016 (comment))

Checklist

• Code follows style guide (ruff)
• Tests added
• Documented the changes

Examples

This PR gives us static typing for Implementation at the narwhals-level:

from __future__ import annotations

from typing import Any

import pandas as pd
import polars as pl
import pyarrow as pa
from typing_extensions import reveal_type

import narwhals as nw

data: dict[str, Any] = {"a": [1, 2, 3]}

impl_1 = nw.from_native(pd.DataFrame(data)).implementation
impl_2 = nw.from_native(pl.DataFrame(data)).implementation
impl_3 = nw.from_native(pa.table(data)).implementation

reveal_type(impl_1)  # Type of "impl_1" is "Literal[Implementation.PANDAS]"
reveal_type(impl_2)  # Type of "impl_2" is "Literal[Implementation.POLARS]"
reveal_type(impl_3)  # Type of "impl_3" is "Literal[Implementation.PYARROW]"

---

Adapting (https://github.com/narwhals-dev/narwhals#example), the fun stuff happens when we use aliases/type variables.

IntoFrameT expands to all valid Implementation's:

if TYPE_CHECKING:
    from narwhals.typing import IntoFrameT


def agnostic_function(df_native: IntoFrameT) -> IntoFrameT:
    frame = nw.from_native(df_native)
    reveal_type(frame.implementation)  # Type of "frame.implementation" is "Literal[Implementation.PANDAS, Implementation.CUDF, Implementation.MODIN, Implementation.PYARROW, Implementation.POLARS, Implementation.PYSPARK, Implementation.SQLFRAME, Implementation.PYSPARK_CONNECT, Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS]"
    return frame.to_native()

However, if we tweak that to use IntoDataFrameT it can be narrowed to all of EagerAllowed:

if TYPE_CHECKING:
    from narwhals.typing import IntoDataFrameT


def agnostic_function(df_native: IntoDataFrameT) -> IntoDataFrameT:
    frame = nw.from_native(df_native)
    reveal_type(frame.implementation)  # Type of "frame.implementation" is "Literal[Implementation.PANDAS, Implementation.CUDF, Implementation.MODIN, Implementation.PYARROW, Implementation.POLARS]"
    return frame.to_native()

The inverse also works, so when we use IntoLazyFrameT it can be narrowed to all of LazyAllowed:

if TYPE_CHECKING:
    from narwhals.typing import IntoLazyFrameT


def agnostic_function(df_native: IntoLazyFrameT) -> IntoLazyFrameT:
    frame = nw.from_native(df_native)
    reveal_type(frame.implementation)  # Type of "frame.implementation" is "Literal[Implementation.PYSPARK, Implementation.SQLFRAME, Implementation.PYSPARK_CONNECT, Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS, Implementation.POLARS]"
    return frame.to_native()

Note

See tests/implementation_test.py for more exhaustive examples 😅

[FBruzzesi]

I didn't take a look at this yet, but I want to add some context if someone else is reviewing it and wonders what the hell is happening.

In #2983 I am creating a new series via a snippet such as:

x: SeriesT
impl = x.implementation
new_series(..., backend=impl)

Since in #3002 we too good, now the type checker complains since backend is supposed to be a eager Implementation (how can we blame it after all?), however the return type of Series.implementation is just Implementation.

I then started a series of adjustments which led to nowhere, so I asked for help and here is @dangotbanned making magic

[FBruzzesi]

Thanks @dangotbanned - I left a few comments/questions, the main one being: I would already be happy to have the distinction between eager and lazy, so that we know that a DataFrame and a Series always returns a eager implementation (although a lazyframe might return any implementation in practice, since class EagerDataFrame(CompliantDataFrame[...], CompliantLazyFrame[...], ...)

[FBruzzesi]

I am quite happy with this! Thanks @dangotbanned

@MarcoGorelli do you have any objection and/or strong opinion against it?

[FBruzzesi]

👌🏼

[dangotbanned]

I am quite happy with this! Thanks @dangotbanned

Thanks @FBruzzesi!

@MarcoGorelli do you have any objection and/or strong opinion against it?

I cannot stress this enough ...

This PR is easily the most complex bit of typing I've ever worked on 😳
I hope that all the docs and tests help break down the chunkier parts

But I had no idea this was possible a week ago when I started it 😂

[MarcoGorelli]

just one comment, else looks good, thanks!