[Bug]: Typing--cannot chain function taking FrameT with nw.from_native()

[CangyuanLi]

Describe the bug

Hi, not sure if this is a bug or just me using the type system wrong, but I am running into typing errors when I pass in the result of nw.from_native() (when the input is an nwt.IntoFrame) to a function that accepts a FrameT and outputs a FrameT.

Steps or code to reproduce the bug

from typing import reveal_type

import narwhals as nw
import narwhals.typing as nwt


def nw_in_nw_out(df: nwt.FrameT) -> nwt.FrameT:
    return df


def any_df_in_float_out(df: nwt.IntoFrame) -> float:
    nw_df = nw.from_native(df)
    x = nw_in_nw_out(nw_df)
    # Argument of type "DataFrame[IntoDataFrame] | LazyFrame[NativeLazyFrame]
    # | LazyFrame[NativeIbis]" cannot be assigned to parameter "df" of type "FrameT@nw_in_nw_out" in function "nw_in_nw_out"
    # Type "DataFrame[IntoDataFrame] | LazyFrame[NativeLazyFrame] | LazyFrame[NativeIbis]"
    # is not assignable to constrained type variable "FrameT"

    if isinstance(x, nw.LazyFrame):
        x = x.collect()

    return x["x"].mean()


def any_df_in_float_out2(df: nwt.IntoFrame) -> float:
    nw_df = nw.from_native(df)
    nw_df: nwt.Frame
    x = nw_in_nw_out(nw_df)

    reveal_type(x)  # type of x is DataFrame[Any]-- shouldn't lazyframes be OK as well?

    if isinstance(x, nw.LazyFrame):
        x = x.collect()

    return x["x"].mean()

Expected results

I would expect that no typing error be thrown.

Actual results

In the any_df_in_float_out function, I think I need to use nwt.IntoFrame because I accept any convertible DataFrame or LazyFrame, but return a python float (so nwt.IntoFrameT is not appropriate). In the nw_in_nw_out function, I take in a FrameT and output a FrameT, which I think is correct--"Use this if your function accepts either nw.DataFrame or nw.LazyFrame and returns an object of the same kind." However, this results in

# Argument of type "DataFrame[IntoDataFrame] | LazyFrame[NativeLazyFrame]
# | LazyFrame[NativeIbis]" cannot be assigned to parameter "df" of type "FrameT@nw_in_nw_out" in function "nw_in_nw_out"
# Type "DataFrame[IntoDataFrame] | LazyFrame[NativeLazyFrame] | LazyFrame[NativeIbis]"
# is not assignable to constrained type variable "FrameT"

One workaround I found was to explicitly type nw_df as an nwt.Frame. However, this makes the result of nw_in_nw_out to be a DataFrame[Any], instead of a DataFrame[Any] | LazyFrame[Any]

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

System:
    python: 3.12.9 (tags/v3.12.9:fdb8142, Feb  4 2025, 15:27:58) [MSC v.1942 64 bit (AMD64)]
executable: C:\Users\E470032\OneDrive - Fifth Third Bancorp\Documents\data-services\.venv\Scripts\python.exe
   machine: Windows-11-10.0.22631-SP0

Python dependencies:
     narwhals: 2.16.0
        numpy: 2.4.1
       pandas: 2.3.3
        modin:
         cudf:
      pyarrow: 23.0.0
      pyspark:
       polars: 1.37.1
         dask:
       duckdb:
         ibis: 11.0.0
     sqlframe:

Relevant log output