[Bug]: `TypeVar`(s) used in `nw.(BaseFrame|Series)` are recursive

[dangotbanned]

Describe the bug

Discovered during (#2236 (comment)) that we probably need to narrow the TypeVar(s) used for the narwhals (not compliant) types.

Screenshot

Steps or code to reproduce the bug

import polars as pl
from typing_extensions import reveal_type

import narwhals as nw

series = pl.Series([1, 2, 3])
nw_series = nw.from_native(series, series_only=True)
reveal_type(nw_series)
# Type of "nw_series" is "narwhals.series.Series[polars.series.series.Series]"

nw_series_nest_1 = nw.from_native(nw_series, series_only=True)
reveal_type(nw_series_nest_1)
# Type of "nw_series_nest_1" is "narwhals.series.Series[narwhals.series.Series[polars.series.series.Series]]"

nw_series_nest_2 = nw.from_native(nw_series_nest_1, series_only=True)
reveal_type(nw_series_nest_2)
# Type of "nw_series_nest_2" is "narwhals.series.Series[narwhals.series.Series[narwhals.series.Series[polars.series.series.Series]]]"

Expected results

The nested types shouldn't be valid and are currently inaccurate.

At runtime we have:

>>> type(nw_series_nest_1._compliant_series)
narwhals._polars.series.PolarsSeries

>>> type(nw_series_nest_2._compliant_series)
narwhals._polars.series.PolarsSeries

Actual results

nw.from_native is masking what's going on, so a simple example is to try and create that type manually.

import polars as pl

from narwhals._polars.series import PolarsSeries
from narwhals.utils import Version
from narwhals.utils import parse_version

compliant_series = PolarsSeries(
    series, backend_version=parse_version(pl), version=Version.MAIN
)
nw_series_manual = nw.Series(compliant_series, level="full")
reveal_type(nw_series_manual)
# Type of "nw_series_manual" is "Series[Unknown]"

So far, that's expected until this is resolved in #2236:

narwhals/narwhals/series.py

Lines 80 to 82 in 275c5b6

	# TODO @dangotbanned: Repeat (#2119) for `CompliantSeries` to support typing
	# morally: `CompliantSeries`
	self._compliant_series = series.__narwhals_series__()

But now we can force the error:

>>> nw.Series(nw_series_manual, level="full")
AssertionError: Expected Polars Series or an object which implements `__narwhals_series__`, got: <class 'narwhals.series.Series'>.

Please run narwhals.show_version() and enter the output below.

275c5b6

[dangotbanned]

Solution

I've tried to keep this self-contained.
I imagine fixing this is going to be more work.

Generally, it'll be applying this idea - in combination with #2116

from __future__ import annotations

from typing import Generic
from typing import TypeVar

import polars as pl
from typing_extensions import reveal_type

from narwhals._polars.series import PolarsSeries
from narwhals.utils import Version
from narwhals.utils import parse_version

if TYPE_CHECKING:
    from narwhals.typing import NativeSeries
    from narwhals.utils import _StoresNative


SeriesT = TypeVar("SeriesT", bound="NativeSeries")
FromSeriesT = TypeVar("FromSeriesT", bound="NativeSeries")
"""**Very important** that this isn't the same `TypeVar` as `__init__`."""


class Series2(Generic[SeriesT]):
    _series: SeriesT

    # - In reality, this'll be `_StoresNative[SeriesT]`
    # - But `Series2` will still be generic over `SeriesT` only
    def __init__(self, series: SeriesT) -> None: 
        self._series = series

    def to_native(self) -> SeriesT:
        return self._series

    @classmethod
    def from_native(cls, series: FromSeriesT) -> Series2[FromSeriesT]:
        return Series2(series)

    @classmethod
    def from_narwhals(cls, series: Series2[FromSeriesT]) -> Series2[FromSeriesT]:
        return series

    @classmethod
    def from_compliant(cls, series: _StoresNative[FromSeriesT]) -> Series2[FromSeriesT]:
        return Series2(series.native)


pl_series: pl.Series = pl.Series([1, 2, 3])
compliant_series = PolarsSeries(
    series, backend_version=parse_version(pl), version=Version.MAIN
)

okay = Series2(pl_series)
reveal_type(okay)
# Type of "okay" is "Series2[Series]"

not_okay = Series2(Series2(pl_series))
reveal_type(not_okay)
# Type of "not_okay" is "Series2[Unknown]"

okay_nw = Series2.from_narwhals(okay)
reveal_type(okay_nw)
# Type of "okay_nw" is "Series2[Series]"

okay_native = Series2.from_native(pl_series)
reveal_type(okay_native)
# Type of "okay_native" is "Series2[Series]"

okay_compliant = Series2.from_compliant(compliant_series)
reveal_type(okay_compliant)
# Type of "okay_compliant" is "Series2[Series]"

Red squiggles proof

Notes

• from_narwhals and from_compliant can be made private
• They're included to try and define all of the interactions we need for (refactor: adds _Stores(Native|Compliant) #2130 (comment))

[dangotbanned]

A similar issue is blocking (#2232) from being able to annotate LazyFrame._compliant_frame.

FrameT doesn't have a clear relationship with CompliantLazyFrame

Screenshot

The typing here seems quite off:

narwhals/narwhals/dataframe.py

Line 67 in 275c5b6

FrameT = TypeVar("FrameT", bound="IntoFrame")

narwhals/narwhals/typing.py

Lines 75 to 78 in 275c5b6

	IntoFrame: TypeAlias = Union[
	"NativeFrame", "DataFrame[Any]", "LazyFrame[Any]", "DataFrameLike"
	]
	"""Anything which can be converted to a Narwhals DataFrame or LazyFrame.

[dangotbanned]

@MarcoGorelli pinging in case you missed this one
Let me know if there is anything that isn't clear 👍

I'd like to decide how we should resolve prior to v2, as this has the potential for a lot of churn if we aren't careful.

Although, backcompat/#exceptions may have carved out space for us 🤔

• We may consider making some type hints more precise.

I imagine the bug will become a lot more visible after (#2116), where we aren't hiding behind nw.from_native anymore

[MarcoGorelli]

thanks for the ping - saw this but wasn't sure what to suggest 😳

[dangotbanned]

thanks for the ping - saw this but wasn't sure what to suggest 😳

@MarcoGorelli

Ah I see 😄

Just so we're on the same page I'll try to summarize.

Example

Using nw.Series as an example because it is simpler than nw.(Data|Lazy)Frame

I'm trying to say we need to make this change:

Current

from __future__ import annotations

from typing import TYPE_CHECKING
from typing import Any
from typing import TypeVar

if TYPE_CHECKING:
    from typing_extensions import TypeAlias

    from narwhals._typing import CompliantSeries
    from narwhals._typing import NativeSeries

IntoSeries: TypeAlias = "Series[Any] | NativeSeries"  # <------------- Generic over itself
IntoSeriesT = TypeVar("IntoSeriesT", bound=IntoSeries)

class Series(Generic[IntoSeriesT]):
    self._compliant_series: CompliantSeries[IntoSeriesT]

Proposed

from __future__ import annotations

from typing import TYPE_CHECKING
from typing import TypeVar

if TYPE_CHECKING:
    from narwhals._typing import CompliantSeries
    from narwhals._typing import NativeSeries

SeriesT = TypeVar("SeriesT", bound="NativeSeries")  # <------------- Generic over the same type as `Compliant*`

class Series(Generic[SeriesT]):
    self._compliant_series: CompliantSeries[SeriesT]

Issue(s)

• nw.(Data|Lazy)Frame
  • also has to deal with DataFrameLike
  • .lazy().collect() seems it might complicate things
• Most of the TypeVar(s) and TypeAlias(s) in nw.typing are impacted
• Changing main will also change v1

[MarcoGorelli]

sure, happy to explore this, and as you noted, we indeed do have a "We may consider making some type hints more precise." escape hatch

[dangotbanned]

Great thanks @MarcoGorelli, I'll take another look as soon as it becomes an issue for #2116