Feature request: support jaxtyping type hints (currently valid type hints flagged incorrectly)

[jaanli]

Minimal example that incorrectly flags a valid type hint using the jaxtyping type hint library (https://github.com/patrick-kidger/jaxtyping):

import numpy as np
from jaxtyping import Int
from beartype import beartype

@beartype
def predict_labels_for_encounter(
    encounter_index: np.int32,
    val_sentence_neighbor_idx: Int[np.ndarray, "num_nearest_neighbors"]):
    pass

[crypdick]

I found a work-around to this problem by changing Ruff's typing modules. However, it would be awesome if Ruff would ship this alias within the Ruff codebase. cc @charliermarsh @leycec

[charliermarsh]

We may be able to add some of these as first-class members. I need to look into the relationship between jaxtyping and beartype though.

[leycec]

@beartype lead @leycec is summoned. Thanks so much for the kind consideration, @charliermarsh. Likewise, thanks for pinging me here, @crypdick. Let's also ping jaxtyping lead @patrick-kidger. This is his baby.

I need to look into the relationship between jaxtyping and beartype though.

Tenuous at best. @patrick-kidger could probably fill us in on all the details. Briefly:

• jaxtyping type-checks jaxtyping-specific type hints (e.g., jaxtyping.Int[...] in the above example).
• jaxtyping then loosely couples with a more general-purpose type-checker to type-check the remaining jaxtyping-agnostic type hints (e.g., numpy.int32 in the above example). Because the coupling is very loose indeed, jaxtyping literally supports all possible general-purpose type-checkers satisfying the standard decorator protocol. This includes:
  • @beartype, of course. Yay! :partying_face:
  • typeguard, of course. Boo! Kidding. Only kidding. Everybody loves typeguard here. :eyes:

Honestly, this issue is probably just about @beartype. As awesome as jaxtyping and typeguard are, this really has nothing to do with either of those packages – I think, anyway. Let us know if @beartype can do anything to assist Ruff here.

[patrick-kidger]

Hey folks! jaxtyping is my package. This is a known issue, and I recommend just disabling the F722 check on this.

(@leycec - this issue is purely jaxtyping-specific, and unrelated to beartype, you'll be pleased to know!)

What's happening is that in type annotations, strings are either (a) just strings or (b) forward references. Ruff is assuming that only the second case is true, trying to resolve our strings as forward references, and failing. This is consistent with flake8.

To be fair, Python itself offers no clear way of distinguishing these cases either. Most static type checkers have to special-case Annotated so that Annotated[Foo, "bar"] is still valid. We actually use if TYPE_CHECKING: from typing import Annotated as Float to make things work with them, but it's not clear if tools like ruff should respect TYPE_CHECKING themselves.

If you'd like to special-case ruff support for jaxtyping that would be awesome. Conversely I recognise you may not wish to add special support to a third-party package, in which case I think no action is needed from ruff. :)