typing: `mypy` infers `from_native(IntoFrameT) -> IntoDataFrameT`

[afrisgaard]

What type of report is this?

Correction

Please describe the issue.

My linter is not happy when I use the combined typeVar IntoFrameT for both DataFrame and LazyFrame with from_native()

It seems it's only happy with me using the IntoDataFrameT or IntoLazyFrameT and not the IntoFrameT notation.

Maybe I'm using it wrong. But from the docs it seems like functions where I handle both Lazy and Eager cases, should be IntoFrameT.

Using narwhals 2.12.0.
Happy:

Un-happy:

If you have a suggestion on how it should be, add it below.

No response

[dangotbanned]

Could you share a full (code) repro I can copy to check out please?

I want to try this out with (#3125) as well

[afrisgaard]

Hey!

Did a very naive example here:
https://github.com/afrisgaard/tryout/blob/main/main.py

Should be public

[dangotbanned]

Thanks @afrisgaard, I've just put this together to show a bit more about what is going on.

Repro

from __future__ import annotations  # narwhals/t.py

from typing import TYPE_CHECKING

from typing_extensions import reveal_type

import narwhals as nw

if TYPE_CHECKING:
    from narwhals.typing import IntoFrameT


def repro_3319_1(df_in: IntoFrameT) -> IntoFrameT:
    reveal_type(df_in)

    df = nw.from_native(df_in)
    reveal_type(df)

    result = df.to_native()
    reveal_type(result)
    return result


def repro_3319_2(df_in: IntoFrameT) -> IntoFrameT:
    return nw.from_native(df_in).to_native()


def api(frame: IntoFrameT) -> tuple[IntoFrameT, IntoFrameT, IntoFrameT]:
    reveal_type(frame)

    result_1 = repro_3319_1(frame)
    reveal_type(result_1)

    result_2 = repro_3319_2(frame)
    reveal_type(result_2)
    return frame, result_1, result_2

Show mypy output 1

Config

narwhals/t.py:14: note: Revealed type is "IntoFrameT`-1"
narwhals/t.py:16: error: Value of type variable "IntoLazyFrameT" of "from_native" cannot be "IntoFrameT"  [type-var]
        df = nw.from_native(df_in)
             ^~~~~~~~~~~~~~~~~~~~~
narwhals/t.py:17: note: Revealed type is "narwhals.dataframe.LazyFrame[IntoFrameT`-1]"
narwhals/t.py:20: note: Revealed type is "IntoFrameT`-1"
narwhals/t.py:25: error: Value of type variable "IntoLazyFrameT" of "from_native" cannot be "IntoFrameT"  [type-var]
        return nw.from_native(df_in).to_native()
               ^~~~~~~~~~~~~~~~~~~~~
narwhals/t.py:29: note: Revealed type is "IntoFrameT`-1"
narwhals/t.py:32: note: Revealed type is "IntoFrameT`-1"
narwhals/t.py:35: note: Revealed type is "IntoFrameT`-1"
Found 2 errors in 1 file (checked 1 source file)

Show pyright output 1

Config

narwhals/t.py
  narwhals/t.py:14:17 - information: Type of "df_in" is "IntoFrameT@repro_3319_1"
  narwhals/t.py:17:17 - information: Type of "df" is "DataFrame[NativeDataFrame*]* | LazyFrame[NativeLazyFrame*]* | LazyFrame[NativeIbis*]*"
  narwhals/t.py:20:17 - information: Type of "result" is "NativeDataFrame* | NativeIbis* | NativeLazyFrame*"
  narwhals/t.py:29:17 - information: Type of "frame" is "IntoFrameT@api"
  narwhals/t.py:32:17 - information: Type of "result_1" is "IntoFrameT@api"
  narwhals/t.py:35:17 - information: Type of "result_2" is "IntoFrameT@api"
0 errors, 0 warnings, 6 informations

But what does any of that mean?

If we extend the repro to use the type pl.DataFrame | pl.LazyFrame to solve IntoFrameT and keep going:

Show much longer repro

from __future__ import annotations  # narwhals/t.py

from typing import TYPE_CHECKING

from typing_extensions import reveal_type

import narwhals as nw

if TYPE_CHECKING:
    import polars as pl

    from narwhals.typing import IntoFrameT


def repro_3319_1(df_in: IntoFrameT) -> IntoFrameT:
    reveal_type(df_in)

    nw.from_native(df_in)

    df = nw.from_native(df_in)
    reveal_type(df)

    result = df.to_native()
    reveal_type(result)
    return result


def repro_3319_2(df_in: IntoFrameT) -> IntoFrameT:
    return nw.from_native(df_in).to_native()


def api(frame: IntoFrameT) -> tuple[IntoFrameT, IntoFrameT, IntoFrameT]:
    reveal_type(frame)

    result_1 = repro_3319_1(frame)
    reveal_type(result_1)

    result_2 = repro_3319_2(frame)
    reveal_type(result_2)
    return frame, result_1, result_2


def no_type_var(frame: pl.DataFrame | pl.LazyFrame) -> None:
    reveal_type(frame)

    result_1 = repro_3319_1(frame)
    reveal_type(result_1)

    result_2 = repro_3319_2(frame)
    reveal_type(result_2)

    result_3 = api(frame)
    reveal_type(result_3)

    nw_frame = nw.from_native(frame)
    reveal_type(nw_frame)

    original_frame = nw_frame.to_native()
    reveal_type(original_frame)

    okay = nw_frame.select("a", "b", "c").with_columns(a_max=nw.col("a").max())
    reveal_type(okay)

    bad = nw_frame.to_polars()
    reveal_type(bad)

Show mypy output 2

...
narwhals/t.py:44: note: Revealed type is "Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame]"
narwhals/t.py:47: note: Revealed type is "Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame]"
narwhals/t.py:50: note: Revealed type is "Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame]"
narwhals/t.py:53: note: Revealed type is "tuple[Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame], Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame], Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame]]"
narwhals/t.py:56: note: Revealed type is "Union[narwhals.dataframe.DataFrame[polars.dataframe.frame.DataFrame], narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]]"
narwhals/t.py:59: note: Revealed type is "Union[polars.dataframe.frame.DataFrame, polars.lazyframe.frame.LazyFrame]"
narwhals/t.py:62: note: Revealed type is "Union[narwhals.dataframe.DataFrame[polars.dataframe.frame.DataFrame], narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]]"
narwhals/t.py:64: error: Item "narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]" of
"narwhals.dataframe.DataFrame[polars.dataframe.frame.DataFrame] | narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]" has no attribute "to_polars"  [union-attr]
        bad = nw_frame.to_polars()
              ^~~~~~~~~~~~~~~~~~
narwhals/t.py:65: note: Revealed type is "Union[polars.dataframe.frame.DataFrame, Any]"
Found 4 errors in 1 file (checked 1 source file)

Show pyright output 2

narwhals/t.py
  ...
  narwhals/t.py:44:17 - information: Type of "frame" is "DataFrame | LazyFrame"
  narwhals/t.py:47:17 - information: Type of "result_1" is "DataFrame | LazyFrame"
  narwhals/t.py:50:17 - information: Type of "result_2" is "DataFrame | LazyFrame"
  narwhals/t.py:53:17 - information: Type of "result_3" is "tuple[DataFrame | LazyFrame, DataFrame | LazyFrame, DataFrame | LazyFrame]"
  narwhals/t.py:56:17 - information: Type of "nw_frame" is "narwhals.dataframe.DataFrame[polars.dataframe.frame.DataFrame] | narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]"
  narwhals/t.py:59:17 - information: Type of "original_frame" is "DataFrame | LazyFrame"
  narwhals/t.py:62:17 - information: Type of "okay" is "narwhals.dataframe.DataFrame[polars.dataframe.frame.DataFrame] | narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]"    
  narwhals/t.py:64:20 - error: Cannot access attribute "to_polars" for class "narwhals.dataframe.LazyFrame[polars.lazyframe.frame.LazyFrame]"
    Attribute "to_polars" is unknown (reportAttributeAccessIssue)
  narwhals/t.py:65:17 - information: Type of "bad" is "DataFrame | Unknown"
1 error, 0 warnings, 14 informations

The error at the very end (for both type checkers) is probably the thing the most users are looking to avoid.

AFAICT, the other differences are because mypy picks the first possible @overload for IntoFrameT.
That causes a problem if it doesn't return DataFrame | LazyFrame, which the case for main (DataFrame), but on #3125 mypy seems to pick the opposite (LazyFrame) 🫠

I think that what pyright is doing is more useful, because it infers (* symbol) that more than one @overload can match

narwhals/t.py:17:17 - information: Type of "df" is "DataFrame[NativeDataFrame*]* | LazyFrame[NativeLazyFrame*]* | LazyFrame[NativeIbis*]*"

Important

The difficulty on our end is that IntoFrameT was introduced before our typing was able to distinguish DataFrame and LazyFrame in from_native - without the use of keyword arguments (#2944)

Suggestion

If you need to use mypy and are writing functions which are truly compatible with both DataFrame and LazyFrame

Option 1

Force mypy to treat df_in as either eager or lazy.

The main downside is you now have a Union to deal with

from __future__ import annotations

from typing import TYPE_CHECKING

from typing_extensions import reveal_type

import narwhals as nw

if TYPE_CHECKING:
    from narwhals.typing import IntoDataFrameT, IntoLazyFrameT

def repro_3319_fix_1(
    df_in: IntoDataFrameT | IntoLazyFrameT,
) -> IntoDataFrameT | IntoLazyFrameT:
    reveal_type(df_in)

    nw.from_native(df_in)

    df = nw.from_native(df_in)
    reveal_type(df)

    result = df.to_native()
    reveal_type(result)
    return result

Option 2

This "resolves" the downside of Option 1, but is pretty verbose and probably not something you want to maintain in multiple places

@overload
def repro_3319_fix_2(df_in: IntoDataFrameT) -> IntoDataFrameT: ...


@overload
def repro_3319_fix_2(df_in: IntoLazyFrameT) -> IntoLazyFrameT: ...


@overload
def repro_3319_fix_2(
    df_in: IntoDataFrameT | IntoLazyFrameT,
) -> IntoDataFrameT | IntoLazyFrameT: ...


def repro_3319_fix_2(
    df_in: IntoDataFrameT | IntoLazyFrameT,
) -> IntoDataFrameT | IntoLazyFrameT:
    reveal_type(df_in)

    nw.from_native(df_in)

    df = nw.from_native(df_in)
    reveal_type(df)

    result = df.to_native()
    reveal_type(result)
    return result

Note

I'm not really satisfied with either option FWIW