How to narrowing a dataclass return type

[jackjyq]

from dataclasses import dataclass


@dataclass
class Row:
    cell: str


@dataclass
class Table:
    row: Row | None


def query(key: int) -> Table | None:
    """
    The Table.row is optional somewhere else, however, in 
    this function, row is guaranteed to be not None when Table is not None.

    Question: how to narrow the type of Table.row in this function?
    """
    if key == 1:
        return Table(
            row=Row(cell="value"),
        )


if table := query(1):
    # there is `"cell" is not a known member of "None"` warning
    print(table.row.cell)

The Table.row is optional somewhere else, however, in this function, row is guaranteed to be not None when Table is not None.

Question: how to narrow the type of Table.row in this function?

[erictraut]

If you're sure that table.row is not None, you can add an assert statement to this effect. This documents your assumption for yourself, other readers of your code, and for static type checkers.

if table := query(1):
    assert table.row is not None
    print(table.row.cell)

[jackjyq]

Hi @erictraut ,

Thanks for your answer.

I think the table.row is not None condition is part of the function, so it is better to narrow the result type within the function, rather than outside.

I have tried to move the assertion in the function, but it does not work. I am wondering is it possible to do this in the query function?

def query(key: int) -> Table | None:
    if key == 1:
        table=  Table(
            row=Row(cell="value"),
        )
        assert table.row is not None
        return table


if table := query(1):
    # there is `"cell" is not a known member of "None"` warning
    print(table.row.cell)

[erictraut]

I am wondering is it possible to do this in the query function?

No, this is not possible to do in another function. The return type of query is declared to be Table | None regardless of the implementation details within that function. Type narrowing for expressions within any execution scope (a module, function, lambda, or comprehension) is local to that execution scope and not observed outside of the scope. This makes sense if you consider that narrowing is performed based on code flow analysis, and code flow graphs are valid only within an execution scope.

There's also no way in the type system to express the type "a Table whose attribute row has been verified not to be None" unless Table is a generic class and parameterized by a type that defines the row attribute. If you were to model it that way, you could specify a return type of Table[Row] rather than Table[Row | None].

[jackjyq]

I see, thanks Eric.