Type narrowing based on variable

[Andarius]

Hi,

I'm trying to build a query builder that would allow the following:

insert_one('my-table', {'foo': 1})
    .returning(key='foo')
    .on_conflict(keys=['id'])

I'm trying to prevent the use of multiple .returning or .conflict if it has already been used.

For now I have been playing around with the type system to see if it was doable with something like this:

from dataclasses import make_dataclass, dataclass
from typing import overload, Literal, Protocol


@dataclass
class QueryBase:
    query: str | None = None


class WithConflict(Protocol):
    def on_conflict(self) -> str:
        ...


class WithReturning(Protocol):
    def returning(self) -> str:
        ...


class WithReturningConflict(WithReturning, WithConflict):
    def on_conflict(self) -> str:
        ...

    def returning(self) -> str:
        ...


@overload
def get_class(with_conflict: Literal[True], with_returning: Literal[True]) -> WithReturningConflict: ...


@overload
def get_class(with_conflict: Literal[True], with_returning: Literal[False]) -> WithConflict: ...


@overload
def get_class(with_conflict: Literal[False], with_returning: Literal[True]) -> WithReturning: ...


@overload
def get_class(with_conflict: Literal[False], with_returning: Literal[False]) -> QueryBase: ...


def get_class(with_conflict: bool = False,
                      with_returning: bool = False
              ):
    _namespace = {}
    if with_conflict:
        _namespace['on_conflict'] = lambda self: 'Conflict'
    if with_returning:
        _namespace['returning'] = lambda self: 'Returning'
    return make_dataclass('C',
                          [],
                          namespace=_namespace,
                          bases=(QueryBase,))()

                          
c2 = get_class(with_conflict=False, with_returning=True)
c2.returning() # Ok
c2.on_conflict() # Error as expected

However as soon as I try to pass it as a variable I fail to make it work:

class Foo(TypedDict):
    returning: Literal[True, False]
    conflict: Literal[True, False]


foo: Foo = {'returning': True, 'conflict': False}

# Pyright throws an error:  Argument of type "bool" cannot be assigned to parameter "with_conflict" of type "Literal[False]" in function "get_class"
c2 = get_class(with_conflict=foo['conflict'], with_returning=foo['returning'])

Is is something that can be done ?

[erictraut]

The type Literal[True, False] is the same as bool.

In your example, you need to provide a fallback overload with bool parameter types. There will be times (like in your final code snippet) where a type checker cannot determine statically whether an argument type is Literal[True] or Literal[False] and can only determine that it's bool. You need an overload that handles this case.

@overload
def get_class(with_conflict: bool, with_returning: bool) -> WithReturningConflict | WithConflict | WithReturning | QueryBase: ...

If you are designing an new API and want it to work well with static type checking, you might want to avoid creating such a polymorphic method — one that returns many varying types.

[Andarius]

Hey @erictraut, thanks for taking the time to answer.

Yeah I'm not a fan either, do you think I could still do it chaining functions without having to use polymorphic methods or should I switch to regular functions with arguments (keeping the main focus on typing here) ?