TypeAlias cannot be found in type_annotation_map

[RazerM]

Describe the bug

I have a TypeAlias that I'm trying to put in type_annotation_map but SQLAlchemy can't seem to find it.

(my naive understanding was that Json in the example would be looked up directly by hash in the annotation map)

Optional link from https://docs.sqlalchemy.org which documents the behavior that is expected

No response

SQLAlchemy Version in Use

2.0.30

DBAPI (i.e. the database driver)

psycopg2

Database Vendor and Major Version

PostgreSQL 12

Python Version

3.10

Operating system

macOS

To Reproduce

from typing import Annotated, Dict, List, Optional, TypeAlias, Union

from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

JsonPrimitive: TypeAlias = Union[str, int, float, bool, None]
JsonObject: TypeAlias = Dict[str, "Json"]
JsonArray: TypeAlias = List["Json"]
Json: TypeAlias = Union[JsonObject, JsonArray, JsonPrimitive]


class Base(DeclarativeBase):
    type_annotation_map = {
        Json: JSONB(none_as_null=True),
    }


class Model(Base):
    __tablename__ = "model"

    id: Mapped[int] = mapped_column(primary_key=True)
    data: Mapped[Optional[Json]]

Error

Traceback (most recent call last):
  File "/scratch_184.py", line 25, in <module>
    class Model(Base):
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/decl_api.py", line 836, in __init_subclass__
    _as_declarative(cls._sa_registry, cls, cls.__dict__)
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/decl_base.py", line 244, in _as_declarative
    return _MapperConfig.setup_mapping(registry, cls, dict_, None, {})
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/decl_base.py", line 325, in setup_mapping
    return _ClassScanMapperConfig(
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/decl_base.py", line 571, in __init__
    self._extract_mappable_attributes()
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/decl_base.py", line 1541, in _extract_mappable_attributes
    value.declarative_scan(
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/properties.py", line 709, in declarative_scan
    self._init_column_for_annotation(
  File "/venv/lib/python3.10/site-packages/sqlalchemy/orm/properties.py", line 880, in _init_column_for_annotation
    raise sa_exc.ArgumentError(
sqlalchemy.exc.ArgumentError: Could not locate SQLAlchemy Core type for Python type typing.Union[str, int, bool, float, typing.Dict[str, ForwardRef('Json')], typing.List[ForwardRef('Json')]] inside the 'data' attribute Mapped annotation

Additional context

There's still an error even if the TypeAlias is not self-referential:

JsonPrimitive: TypeAlias = Union[str, int, float, bool, None]
JsonObject: TypeAlias = Dict[str, object]
JsonArray: TypeAlias = List[object]
Json: TypeAlias = Union[JsonObject, JsonArray, JsonPrimitive]

but weirdly it works if JsonPrimitive doesn't have None:

JsonPrimitive: TypeAlias = Union[str, int, float, bool]
JsonObject: TypeAlias = Dict[str, object]
JsonArray: TypeAlias = List[object]
Json: TypeAlias = Union[JsonArray, str, int, float]

(both of these types aren't very useful, so I would specify the JSONB explicitly in mapped_column, but I thought it's worth mentioning what I'd tried).

[zzzeek]

we probably cant use Unions as keys, kind of an edge case

[CaselIT]

I think the issue here is this:

JsonPrimitive: TypeAlias = Union[str, int, float, bool, None]
JsonObject: TypeAlias = Dict[str, object]
JsonArray: TypeAlias = List[object]
Json: TypeAlias = Union[JsonObject, JsonArray, JsonPrimitive]

print(Optional[Json] == Json)

This prints True. The repr is also the same

@zzzeek I think the current implementation removes the None from the union when searching for type.
It usually makes sense, but not in this case.