Make static type checking work better with symbolics

[charlesyuan314]

This change adds TypeIs to the signature of is_symbolic, thereby making spurious static type checker errors less likely and enabling us to remove uses of cast and assert.

Motivation

A common pattern in the codebase is as follows:

x: SymbolicInt = ...
if is_symbolic(x):
    raise DecomposeTypeError(...)

# do something that assumes that x: int, such as:
c = "abcde"[x]

Unfortunately, this pattern doesn't play well with static type checkers such as mypy, which we use.
As an example we actually encountered, consider this snippet from #1089 (just prior to commit 23577b2):

if is_symbolic(self.k, self.offset):
    raise DecomposeTypeError(f"Cannot decompose symbolic {self=}")

for start in range(0, self.k, 2 * self.offset):
    for step in range(self.offset):
        # do something

Given the above code, mypy fails with the error:

qualtran/bloqs/arithmetic/sorting.py:168: error: Argument 2 to "range" has incompatible type "int | Expr"; expected "SupportsIndex"  [arg-type]
qualtran/bloqs/arithmetic/sorting.py:169: error: Argument 1 to "range" has incompatible type "int | Expr"; expected "SupportsIndex"  [arg-type]

This overly conservative type checking error occurs because in general, a static type checker cannot infer that the type of the argument of is_symbolic should be narrowed to int from Union[int, Expr] when is_symbolic returns False. Such inference is normally reserved for specific language patterns, such as isinstance:

if not isinstance(x, sympy.Expr):
    raise DecomposeTypeError(...)

# here, mypy does actually infer that x: int
c = "abcde"[x]  # no type error

We do sometimes use isinstance in the codebase for this purpose, but it is less general and we would prefer not to replace is_symbolic with it.

Solution

Fortunately, we can do better and get the benefit of type narrowing using the recently introduced TypeIs type annotation. In brief, by annotating the return type of is_symbolic as TypeIs[sympy.Expr] rather than bool, the static type checker will infer that when is_symbolic(arg: Union[int, Expr]) returns True, then arg: Expr, and otherwise arg: int.

This PR makes the above change. Consequently, it safely removes numerous instances of cast(int, ...), assert [not] isinstance(...), int(...), and similar patterns from the code.

There are some caveats due to limitations of the Python type system:

• The benefit does not extend to variadic arguments to is_symbolic. To get the benefit, it is necessary to rewrite is_symbolic(x, y) as is_symbolic(x) or is_symbolic(y).
• The benefit does not extend to other situations such using an is_symbolic(self) method defined on a class rather than symbolics.types.is_symbolic.
• The benefit extends to HasLength and Shaped. However, in one scenario, it was necessary to explicitly annotate k: SymbolicFloat so that the type checker would infer k: Expr rather than k: Union[HasLength, Shaped, Expr] when is_symbolic(k) returns True.
• The benefit does not yet extend to more deeply nested types, such as narrowing Tuple[SymbolicInt, ...] to Tuple[int, ...].

[mpharrigan]

Very cool. What's the minimum required python version with this feature?

[charlesyuan314]

Very cool. What's the minimum required python version with this feature?

Python 3.8+. TypeIs will be native in Python 3.13 and is backported by the typing_extensions package.

[mpharrigan]

sweeeeeeeet

[mpharrigan]

actually can you figure out the minimum version of typing_extensions required and add it to dev_tools/requirements/deps/runtime.txt with a >= constraint. We transitively depend on it already from a variety of other packages and the current pinned version for the CI is sufficient but whatever version I've got locally isn't new enough, and if you pip install qualtran it will respect the runtime.txt loosey-goosey specs rather than the pinned versions.