From f5f532affd84baf6c82a21e4a024dc27cde0123c Mon Sep 17 00:00:00 2001 From: AivanF Date: Thu, 25 Jun 2026 13:58:54 +0300 Subject: [PATCH 1/2] Modern typing --- slayer/api/server.py | 126 +++--- slayer/async_utils.py | 3 +- slayer/cli.py | 20 +- slayer/client/slayer_client.py | 87 ++-- slayer/core/enums.py | 4 +- slayer/core/errors.py | 12 +- slayer/core/format.py | 9 +- slayer/core/formula.py | 75 ++-- slayer/core/models.py | 118 +++--- slayer/core/query.py | 46 +-- slayer/core/refs.py | 13 +- slayer/dbt/converter.py | 63 ++- slayer/dbt/entities.py | 15 +- slayer/dbt/filters.py | 5 +- slayer/dbt/manifest.py | 20 +- slayer/dbt/models.py | 111 +++-- slayer/dbt/parser.py | 13 +- slayer/dbt/sql_resolver.py | 9 +- slayer/demo/jaffle_shop.py | 20 +- slayer/embeddings/client.py | 34 +- slayer/embeddings/models.py | 4 +- slayer/embeddings/ranker.py | 5 +- slayer/engine/column_dependency.py | 54 +-- slayer/engine/column_expansion.py | 33 +- slayer/engine/enriched.py | 108 ++--- slayer/engine/enrichment.py | 185 ++++----- slayer/engine/ingestion.py | 184 ++++----- slayer/engine/profiling.py | 68 +-- slayer/engine/query_engine.py | 194 ++++----- slayer/engine/schema_drift.py | 386 +++++++++--------- slayer/facade/catalog.py | 95 +++-- slayer/facade/catalog_sql.py | 74 ++-- slayer/facade/info_schema.py | 13 +- slayer/facade/probe_queries.py | 3 +- slayer/facade/rows.py | 6 +- slayer/facade/translator.py | 364 ++++++++--------- slayer/flight/_capture_stub.py | 16 +- slayer/flight/auth.py | 17 +- slayer/flight/cli.py | 5 +- slayer/flight/handlers.py | 17 +- slayer/flight/info_schema.py | 3 +- slayer/flight/probe_queries.py | 3 +- slayer/flight/server.py | 14 +- slayer/flight/types.py | 3 +- slayer/help/__init__.py | 3 +- slayer/mcp/server.py | 256 ++++++------ slayer/memories/models.py | 14 +- slayer/memories/ranker.py | 9 +- slayer/memories/resolver.py | 32 +- slayer/memories/service.py | 23 +- slayer/pg_facade/auth.py | 7 +- slayer/pg_facade/connection.py | 100 ++--- slayer/pg_facade/identity.py | 3 +- slayer/pg_facade/probes.py | 25 +- slayer/pg_facade/protocol.py | 29 +- slayer/pg_facade/server.py | 13 +- slayer/pg_facade/types.py | 10 +- slayer/search/cypher_naive.py | 5 +- slayer/search/graph.py | 10 +- slayer/search/index.py | 61 ++- slayer/search/render.py | 31 +- slayer/search/retriever.py | 29 +- slayer/search/retrievers/bm25.py | 21 +- slayer/search/retrievers/embeddings.py | 55 ++- slayer/search/retrievers/tantivy.py | 21 +- slayer/search/rrf.py | 5 +- slayer/search/service.py | 256 ++++++------ slayer/sql/client.py | 77 ++-- slayer/sql/dialects/__init__.py | 4 +- slayer/sql/dialects/_tier2.py | 13 +- slayer/sql/dialects/base.py | 21 +- slayer/sql/dialects/bigquery.py | 7 +- slayer/sql/dialects/clickhouse.py | 4 +- slayer/sql/dialects/duckdb.py | 3 +- slayer/sql/dialects/mysql.py | 4 +- slayer/sql/dialects/postgres.py | 3 +- slayer/sql/dialects/snowflake.py | 12 +- slayer/sql/dialects/sqlite.py | 26 +- slayer/sql/dialects/tsql.py | 13 +- slayer/sql/engine_factory.py | 3 +- slayer/sql/generator.py | 38 +- slayer/sql/sql_predicate.py | 5 +- slayer/sql/sqlite_introspect.py | 13 +- slayer/storage/base.py | 81 ++-- slayer/storage/join_sync.py | 41 +- slayer/storage/migrations.py | 7 +- slayer/storage/sidecar_embedding_store.py | 37 +- slayer/storage/sqlite_storage.py | 72 ++-- slayer/storage/type_refinement.py | 15 +- slayer/storage/v2_datasource_migration.py | 4 +- slayer/storage/v2_memory_migration.py | 4 +- slayer/storage/v4_migration.py | 11 +- slayer/storage/v5_migration.py | 3 +- slayer/storage/yaml_storage.py | 43 +- tests/conftest.py | 2 +- tests/facade/test_catalog.py | 9 +- tests/facade/test_catalog_flatten_filter.py | 9 +- .../test_catalog_sql_metabase_golden.py | 7 +- tests/flight/conftest.py | 4 +- tests/integration/_pg_serve_helpers.py | 11 +- tests/integration/conftest.py | 13 +- tests/integration/conftest_metabase.py | 39 +- tests/integration/test_in_memory_sqlite.py | 2 +- .../integration/test_ingestion_jaffle_shop.py | 5 +- tests/integration/test_integration_flight.py | 33 +- .../test_integration_flight_pyarrow_client.py | 37 +- .../integration/test_integration_pg_facade.py | 6 +- tests/integration/test_mcp_inspect.py | 6 +- tests/integration/test_metabase_e2e.py | 10 +- tests/pg_facade/test_connection.py | 39 +- tests/pg_facade/test_probes.py | 9 +- tests/test_async_engine_disposal.py | 3 +- tests/test_bm25_retriever.py | 2 +- tests/test_cli_memory_string_ids.py | 2 +- tests/test_cli_search_text_output.py | 5 +- tests/test_client.py | 79 ++-- tests/test_column_dependency.py | 3 +- tests/test_edit_datasource_refresh.py | 2 +- tests/test_embedding_retriever.py | 20 +- tests/test_embeddings_client.py | 48 +-- tests/test_embeddings_storage.py | 8 +- tests/test_engine_profiling.py | 3 +- tests/test_entity_description_plumbing.py | 2 +- tests/test_entity_resolution.py | 2 +- tests/test_entry_points_namespacing.py | 4 +- tests/test_force_clean.py | 5 +- tests/test_idempotent_ingestion.py | 25 +- .../test_ingestion_memory_dangling_cleanup.py | 2 +- tests/test_inspect_model_compact.py | 3 +- tests/test_mcp_server.py | 5 +- tests/test_memories_cli.py | 9 +- tests/test_memories_client.py | 2 +- tests/test_memories_inspect.py | 5 +- tests/test_memories_mcp.py | 7 +- tests/test_memories_rest.py | 2 +- tests/test_memories_storage.py | 2 +- tests/test_memory_resolver_memory_branch.py | 2 +- tests/test_memory_string_ids.py | 2 +- tests/test_models_summary_compact.py | 3 +- tests/test_nested_dag_cross_stage_refs.py | 7 +- tests/test_projection_trim.py | 14 +- tests/test_rest_memory_string_ids.py | 2 +- tests/test_retriever_abc.py | 2 +- .../test_save_memory_description_plumbing.py | 3 +- tests/test_save_memory_id_kwarg.py | 2 +- ...save_memory_linked_entities_memory_refs.py | 2 +- tests/test_search_compact.py | 2 +- tests/test_search_compact_plumbing.py | 3 +- tests/test_search_datasource_filter.py | 2 +- tests/test_search_description_indexed.py | 2 +- tests/test_search_graph.py | 2 +- tests/test_search_index.py | 5 +- tests/test_search_invariance.py | 18 +- tests/test_search_named_entity_surfacing.py | 4 +- tests/test_search_render_unified.py | 38 +- tests/test_search_service.py | 20 +- tests/test_search_service_dispatch.py | 24 +- tests/test_search_service_orchestration.py | 32 +- tests/test_search_surfaces.py | 2 +- tests/test_search_three_channel.py | 26 +- tests/test_search_unified.py | 2 +- tests/test_sidecar_embedding_store.py | 8 +- tests/test_sql_client_in_memory_async.py | 3 +- tests/test_startup_ingest.py | 15 +- tests/test_tantivy_retriever.py | 5 +- tests/test_validate_models.py | 19 +- 166 files changed, 2478 insertions(+), 2523 deletions(-) diff --git a/slayer/api/server.py b/slayer/api/server.py index 9bf869cd..8106f613 100644 --- a/slayer/api/server.py +++ b/slayer/api/server.py @@ -2,7 +2,7 @@ import logging from importlib.metadata import PackageNotFoundError, version as _pkg_version -from typing import Any, Dict, List, Optional, Union +from typing import Any from fastapi import FastAPI, HTTPException from pydantic import BaseModel, ConfigDict, Field @@ -31,34 +31,34 @@ class QueryRequest(BaseModel): # to pass through to SlayerQuery's pre-validate hook. model_config = ConfigDict(extra="allow") - name: Optional[str] = None # Run-by-name: backing query for a query-backed model + name: str | None = None # Run-by-name: backing query for a query-backed model # ``source_model`` accepts a string (stored model name) or a dict # — the dict form is an inline ``ModelExtension`` (``{"source_name": # "", "columns": [...], "joins": [...]}``) or an inline # ``SlayerModel`` (``{"name": "...", "sql_table": "...", "data_source": # "...", "columns": [...]}``). The full polymorphism is handled by # ``SlayerQuery.model_validate`` downstream. - source_model: Optional[Union[str, Dict[str, Any]]] = None + source_model: str | dict[str, Any] | None = None # ``measures`` and ``dimensions`` accept bare strings as a shorthand, # mirroring the Python API: ``"*:count"`` is lifted to # ``{"formula": "*:count"}``, ``"status"`` to ``{"name": "status"}``. # ``SlayerQuery``'s before-validators (``_coerce_measures`` / # ``_coerce_dimensions``) do the actual lifting downstream. - measures: Optional[List[Union[str, Dict[str, Any]]]] = None - dimensions: Optional[List[Union[str, Dict[str, Any]]]] = None - time_dimensions: Optional[List[Dict[str, Any]]] = None - filters: Optional[List[str]] = None - order: Optional[List[Dict[str, Any]]] = None - limit: Optional[int] = None - offset: Optional[int] = None - whole_periods_only: Optional[bool] = None + measures: list[str | dict[str, Any]] | None = None + dimensions: list[str | dict[str, Any]] | None = None + time_dimensions: list[dict[str, Any]] | None = None + filters: list[str] | None = None + order: list[dict[str, Any]] | None = None + limit: int | None = None + offset: int | None = None + whole_periods_only: bool | None = None # DEV-1543: opt out of the dim-only auto-dedup GROUP BY. Default # (``None`` here) keeps the v3 SlayerQuery default (``True``). Set # ``False`` to emit raw rows. - distinct_dimension_values: Optional[bool] = None - dry_run: Optional[bool] = None - explain: Optional[bool] = None - variables: Optional[Dict[str, Any]] = None + distinct_dimension_values: bool | None = None + dry_run: bool | None = None + explain: bool | None = None + variables: dict[str, Any] | None = None class QueryListRequest(BaseModel): @@ -73,39 +73,39 @@ class QueryListRequest(BaseModel): model_config = ConfigDict(extra="forbid") - queries: List[Dict[str, Any]] - variables: Optional[Dict[str, Any]] = None - dry_run: Optional[bool] = None - explain: Optional[bool] = None + queries: list[dict[str, Any]] + variables: dict[str, Any] | None = None + dry_run: bool | None = None + explain: bool | None = None class FieldMetadataResponse(BaseModel): - label: Optional[str] = None - format: Optional[NumberFormat] = None + label: str | None = None + format: NumberFormat | None = None class AttributesResponse(BaseModel): - dimensions: Dict[str, FieldMetadataResponse] = {} - measures: Dict[str, FieldMetadataResponse] = {} + dimensions: dict[str, FieldMetadataResponse] = {} + measures: dict[str, FieldMetadataResponse] = {} class QueryResponse(BaseModel): - data: List[Dict[str, Any]] + data: list[dict[str, Any]] row_count: int - columns: List[str] - sql: Optional[str] = None - attributes: Optional[AttributesResponse] = None + columns: list[str] + sql: str | None = None + attributes: AttributesResponse | None = None class IngestRequest(BaseModel): datasource: str - include_tables: Optional[List[str]] = None - exclude_tables: Optional[List[str]] = None - schema_name: Optional[str] = None + include_tables: list[str] | None = None + exclude_tables: list[str] | None = None + schema_name: str | None = None class ValidateModelsRequest(BaseModel): - data_source: Optional[str] = None + data_source: str | None = None class DatasourcePriorityRequest(BaseModel): @@ -114,7 +114,7 @@ class DatasourcePriorityRequest(BaseModel): shape and FastAPI rejects mistyped payloads with 422 instead of silently coercing them downstream. """ - priority: List[str] = [] + priority: list[str] = [] class SaveMemoryRequest(BaseModel): @@ -139,8 +139,8 @@ class SaveMemoryRequest(BaseModel): learning: str linked_entities: Any - id: Optional[str] = None - description: Optional[str] = None + id: str | None = None + description: str | None = None class SearchRequest(BaseModel): @@ -159,12 +159,12 @@ class SearchRequest(BaseModel): model_config = ConfigDict(extra="forbid") - entities: Optional[List[str]] = None - query: Optional[Any] = None - question: Optional[str] = None - datasource: Optional[str] = None + entities: list[str] | None = None + query: Any | None = None + question: str | None = None + datasource: str | None = None max_results: int = Field(default=10, ge=1) - cypher_filter: Optional[str] = None + cypher_filter: str | None = None compact: bool = True @@ -201,7 +201,7 @@ def create_app( # NOSONAR(S3776) — FastAPI route-handler factory; complexity app.mount("/mcp", mcp_app) @app.get("/health") - async def health() -> Dict[str, str]: + async def health() -> dict[str, str]: return {"status": "ok"} @app.post( @@ -220,7 +220,7 @@ async def health() -> Dict[str, str]: }, ) async def query( - request: Union[QueryRequest, QueryListRequest], + request: QueryRequest | QueryListRequest, ) -> QueryResponse: try: # Multi-stage DAG: body is ``{"queries": [...], "variables": ..., @@ -296,7 +296,7 @@ async def query( ) attrs = result.attributes - def _convert_meta(d: dict) -> Dict[str, FieldMetadataResponse]: + def _convert_meta(d: dict) -> dict[str, FieldMetadataResponse]: return {k: FieldMetadataResponse(label=v.label, format=v.format) for k, v in d.items()} attributes = None @@ -331,8 +331,8 @@ def _convert_meta(d: dict) -> Dict[str, FieldMetadataResponse]: @app.get("/models") async def list_models( - data_source: Optional[str] = None, - ) -> List[Dict[str, Any]]: + data_source: str | None = None, + ) -> list[dict[str, Any]]: identities = await storage._list_all_model_identities() result = [] for ds_name, name in identities: @@ -341,7 +341,7 @@ async def list_models( model = await storage.get_model(name, data_source=ds_name) if model is None or model.hidden: continue - entry: Dict[str, Any] = {"name": name, "data_source": ds_name} + entry: dict[str, Any] = {"name": name, "data_source": ds_name} if model.description: entry["description"] = model.description result.append(entry) @@ -361,8 +361,8 @@ async def list_models( ) async def get_model( name: str, - data_source: Optional[str] = None, - ) -> Dict[str, Any]: + data_source: str | None = None, + ) -> dict[str, Any]: try: model = await storage.get_model(name, data_source=data_source) except AmbiguousModelError as exc: @@ -384,7 +384,7 @@ async def get_model( "/models", responses={400: {"description": "Model failed validation (e.g. user-supplied cache fields on a query-backed model, save-time SQL generation failure)."}}, ) - async def create_model(model: SlayerModel) -> Dict[str, str]: + async def create_model(model: SlayerModel) -> dict[str, str]: # Route through engine.save_model so query-backed models get cache # populated (and user-supplied cache fields are rejected). try: @@ -397,7 +397,7 @@ async def create_model(model: SlayerModel) -> Dict[str, str]: "/models/{name}", responses={400: {"description": "Body name does not match path name, or model failed validation."}}, ) - async def update_model(name: str, model: SlayerModel) -> Dict[str, str]: + async def update_model(name: str, model: SlayerModel) -> dict[str, str]: if model.name != name: raise HTTPException( status_code=400, @@ -423,8 +423,8 @@ async def update_model(name: str, model: SlayerModel) -> Dict[str, str]: ) async def delete_model( name: str, - data_source: Optional[str] = None, - ) -> Dict[str, Any]: + data_source: str | None = None, + ) -> dict[str, Any]: try: deleted = await storage.delete_model(name, data_source=data_source) except AmbiguousModelError as exc: @@ -440,7 +440,7 @@ async def delete_model( return {"status": "deleted", "name": name} @app.get("/datasources/priority") - async def get_datasource_priority() -> Dict[str, List[str]]: + async def get_datasource_priority() -> dict[str, list[str]]: return {"priority": await storage.get_datasource_priority()} @app.put( @@ -454,7 +454,7 @@ async def get_datasource_priority() -> Dict[str, List[str]]: } }, ) - async def put_datasource_priority(body: DatasourcePriorityRequest) -> Dict[str, Any]: + async def put_datasource_priority(body: DatasourcePriorityRequest) -> dict[str, Any]: priority = list(body.priority) try: await storage.set_datasource_priority(priority) @@ -463,18 +463,18 @@ async def put_datasource_priority(body: DatasourcePriorityRequest) -> Dict[str, return {"status": "ok", "priority": priority} @app.get("/datasources") - async def list_datasources() -> List[Dict[str, Any]]: + async def list_datasources() -> list[dict[str, Any]]: result = [] for name in await storage.list_datasources(): ds = await storage.get_datasource(name) - entry: Dict[str, Any] = {"name": name} + entry: dict[str, Any] = {"name": name} if ds: entry["type"] = ds.type result.append(entry) return result @app.get("/datasources/{name}") - async def get_datasource(name: str) -> Dict[str, Any]: + async def get_datasource(name: str) -> dict[str, Any]: ds = await storage.get_datasource(name) if ds is None: raise HTTPException( @@ -488,12 +488,12 @@ async def get_datasource(name: str) -> Dict[str, Any]: return data @app.post("/datasources") - async def create_datasource(datasource: DatasourceConfig) -> Dict[str, str]: + async def create_datasource(datasource: DatasourceConfig) -> dict[str, str]: await storage.save_datasource(datasource) return {"status": "created", "name": datasource.name} @app.delete("/datasources/{name}") - async def delete_datasource(name: str) -> Dict[str, Any]: + async def delete_datasource(name: str) -> dict[str, Any]: deleted = await storage.delete_datasource(name) if not deleted: raise HTTPException( @@ -516,7 +516,7 @@ async def delete_datasource(name: str) -> Dict[str, Any]: ) async def validate_models_endpoint( request: ValidateModelsRequest, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Diff persisted SlayerModels against live DB schemas. Read-only.""" if request.data_source is not None: ds_check = await storage.get_datasource(request.data_source) @@ -555,7 +555,7 @@ async def validate_models_endpoint( }, }, ) - async def ingest(request: IngestRequest) -> Dict[str, Any]: + async def ingest(request: IngestRequest) -> dict[str, Any]: # Strip newlines from the user-controlled datasource name before it # reaches the log or the response detail (S5145 — log-injection # surface). The ds value already round-tripped through Pydantic so @@ -634,7 +634,7 @@ async def ingest(request: IngestRequest) -> Dict[str, Any]: } }, ) - async def save_memory(request: SaveMemoryRequest) -> Dict[str, Any]: + async def save_memory(request: SaveMemoryRequest) -> dict[str, Any]: try: response = await memory_service.save_memory( learning=request.learning, @@ -657,7 +657,7 @@ async def save_memory(request: SaveMemoryRequest) -> Dict[str, Any]: 404: {"description": "Memory not found."}, }, ) - async def delete_memory(memory_id: str) -> Dict[str, Any]: + async def delete_memory(memory_id: str) -> dict[str, Any]: try: response = await memory_service.forget_memory(identifier=memory_id) except MemoryNotFoundError as exc: @@ -683,7 +683,7 @@ async def delete_memory(memory_id: str) -> Dict[str, Any]: } }, ) - async def search(request: SearchRequest) -> Dict[str, Any]: + async def search(request: SearchRequest) -> dict[str, Any]: try: response = await search_service.search( entities=request.entities, diff --git a/slayer/async_utils.py b/slayer/async_utils.py index 383cbfeb..7daf559c 100644 --- a/slayer/async_utils.py +++ b/slayer/async_utils.py @@ -6,7 +6,8 @@ import asyncio import concurrent.futures -from typing import Any, Coroutine, TypeVar +from typing import Any, TypeVar +from collections.abc import Coroutine T = TypeVar("T") diff --git a/slayer/cli.py b/slayer/cli.py index 1fe4de05..e39789f8 100644 --- a/slayer/cli.py +++ b/slayer/cli.py @@ -5,7 +5,7 @@ import json import os import sys -from typing import Any, Dict, List, Optional +from typing import Any from pydantic import BaseModel, Field @@ -68,8 +68,8 @@ class RefreshSamplesResult(BaseModel): that didn't resolve in the requested scope — those are reported as a hard error so typos fail fast.""" - errors: List[str] = Field(default_factory=list) - unresolved_models: List[str] = Field(default_factory=list) + errors: list[str] = Field(default_factory=list) + unresolved_models: list[str] = Field(default_factory=list) def _add_storage_arg(parser): @@ -727,8 +727,8 @@ async def _refresh_samples_async(*, args, storage) -> "RefreshSamplesResult": and (optional) model filters, accumulate per-column errors and the names of any user-specified models that didn't resolve.""" engine = SlayerQueryEngine(storage=storage) - errors: List[str] = [] - unresolved_models: List[str] = [] + errors: list[str] = [] + unresolved_models: list[str] = [] data_source = args.data_source models = args.models if data_source is None: @@ -891,7 +891,7 @@ def _run_storage_migrate_types(args) -> None: def _resolve_datasource_for_cli_refinement( *, inner, ds_name: str, model_name: str, needs_double: bool, -) -> Optional[Any]: +) -> Any | None: """Resolve the datasource for ``slayer storage migrate-types``. Returns the ``DatasourceConfig`` when present, ``None`` when missing @@ -918,7 +918,7 @@ def _resolve_datasource_for_cli_refinement( def _print_refinement_diff( - *, ds_name: str, model_name: str, upgraded: dict, types_before: Dict[str, str], + *, ds_name: str, model_name: str, upgraded: dict, types_before: dict[str, str], ) -> None: """Print before→after diffs for columns whose type changed during refinement. Migration-only aliases are excluded because ``types_before`` @@ -986,7 +986,7 @@ def _refine_one_model_for_cli( return True -async def _load_raw_model_dict(storage, data_source: str, name: str) -> Optional[dict]: +async def _load_raw_model_dict(storage, data_source: str, name: str) -> dict | None: """Read a model's raw on-disk dict bypassing Pydantic's validator chain.""" import json as _json import os as _os @@ -1208,7 +1208,7 @@ def _run_ingest(args): ) -def _format_edit_entry_lines(entry) -> List[str]: +def _format_edit_entry_lines(entry) -> list[str]: lines = [f"EDIT MODEL: {entry.model_name} (datasource: {entry.data_source})"] for attr, label in _REMOVE_SECTIONS: values = getattr(entry.remove, attr) @@ -1223,7 +1223,7 @@ def _format_validate_models_output(entries) -> str: """Render a List[ToDeleteEntry] as human-readable text for CLI output.""" if not entries: return "No drift detected." - lines: List[str] = [] + lines: list[str] = [] for entry in entries: if entry.tool == "delete_model": lines.append( diff --git a/slayer/client/slayer_client.py b/slayer/client/slayer_client.py index fbe17805..2398a3ad 100644 --- a/slayer/client/slayer_client.py +++ b/slayer/client/slayer_client.py @@ -1,16 +1,15 @@ """Python client for SLayer API.""" import logging -from collections.abc import Mapping as ABCMapping, Sequence as ABCSequence +from collections.abc import ( + Mapping, + Mapping as ABCMapping, + Sequence, + Sequence as ABCSequence, +) from typing import ( TYPE_CHECKING, Any, - Dict, - List, - Mapping, - Optional, - Sequence, - Union, ) from urllib.parse import quote @@ -33,12 +32,12 @@ # ``{"queries": [...]}`` (REST) also accept; ``str`` runs a query-backed model # by name. ``Mapping``/``Sequence`` (not ``Dict``/``List``) so callers passing # ``list[dict[str, str]]`` aren't rejected by pyright's invariance check. -QueryInput = Union[ - SlayerQuery, - Mapping[str, Any], - Sequence[Union[SlayerQuery, Mapping[str, Any]]], - str, -] +QueryInput = ( + SlayerQuery + | Mapping[str, Any] + | Sequence[SlayerQuery | Mapping[str, Any]] + | str +) class SlayerClient: @@ -62,7 +61,7 @@ class SlayerClient: def __init__( self, url: str = "http://localhost:5143", - storage: Optional[Any] = None, + storage: Any | None = None, ): self.url = url.rstrip("/") self._storage = storage @@ -76,8 +75,8 @@ async def _request( self, method: str, path: str, - json: Optional[Dict] = None, - params: Optional[Dict] = None, + json: dict | None = None, + params: dict | None = None, ) -> Any: try: import httpx @@ -94,8 +93,8 @@ def _request_sync( self, method: str, path: str, - json: Optional[Dict] = None, - params: Optional[Dict] = None, + json: dict | None = None, + params: dict | None = None, ) -> Any: try: import httpx @@ -109,7 +108,7 @@ def _request_sync( return resp.json() @staticmethod - def _validated_dump(payload: Mapping[str, Any]) -> Dict[str, Any]: + def _validated_dump(payload: Mapping[str, Any]) -> dict[str, Any]: """Round-trip a single-query payload through ``SlayerQuery`` so the server sees the normalised JSON-mode shape. Necessary because the server's ``QueryRequest`` declares ``measures`` / @@ -128,7 +127,7 @@ def _build_query_body( *, dry_run: bool = False, explain: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Convert any accepted input shape into the JSON body for ``POST /query``. Single source of truth shared by sync + async transports. Never mutates caller-owned dicts or lists. @@ -149,7 +148,7 @@ def _build_query_body( both accept them). """ if isinstance(query, str): - body: Dict[str, Any] = {"name": query} + body: dict[str, Any] = {"name": query} elif isinstance(query, SlayerQuery): body = query.model_dump(mode="json", exclude_none=True) elif isinstance(query, ABCSequence) and not isinstance( @@ -157,7 +156,7 @@ def _build_query_body( ): # ``str`` is also a Sequence but already handled above; guard the # binary string types here too so they raise via the else-branch. - serialised: List[Dict[str, Any]] = [] + serialised: list[dict[str, Any]] = [] for i, item in enumerate(query): if isinstance(item, SlayerQuery): serialised.append( @@ -214,7 +213,7 @@ def _parse_response(result: dict) -> SlayerResponse: """Parse an API JSON response into a SlayerResponse.""" from slayer.core.format import NumberFormat - def _parse_meta_dict(d: dict) -> Dict[str, FieldMetadata]: + def _parse_meta_dict(d: dict) -> dict[str, FieldMetadata]: out = {} for k, v in (d or {}).items(): fmt = None @@ -274,7 +273,7 @@ async def explain(self, query: QueryInput) -> SlayerResponse: """ return await self.query(query=query, explain=True) - async def list_models(self, data_source: Optional[str] = None) -> List[str]: + async def list_models(self, data_source: str | None = None) -> list[str]: if self._storage is not None: names = await self._storage.list_models(data_source=data_source) return list(names) @@ -284,29 +283,29 @@ async def list_models(self, data_source: Optional[str] = None) -> List[str]: async def get_model( self, name: str, - data_source: Optional[str] = None, - ) -> Optional[Any]: + data_source: str | None = None, + ) -> Any | None: if self._storage is not None: return await self._storage.get_model(name, data_source=data_source) params = {"data_source": data_source} if data_source else None return await self._request(method="GET", path=f"/models/{name}", params=params) # NOSONAR(S1192) — REST path is the API contract; defining a constant adds indirection without value - async def create_model(self, model: Dict[str, Any]) -> Dict[str, str]: + async def create_model(self, model: dict[str, Any]) -> dict[str, str]: return await self._request(method="POST", path="/models", json=model) # NOSONAR(S1192) — REST path is the API contract; defining a constant adds indirection without value - async def list_datasources(self) -> List[str]: + async def list_datasources(self) -> list[str]: return await self._request(method="GET", path="/datasources") - async def create_datasource(self, datasource: Dict[str, Any]) -> Dict[str, str]: + async def create_datasource(self, datasource: dict[str, Any]) -> dict[str, str]: return await self._request(method="POST", path="/datasources", json=datasource) - async def get_datasource_priority(self) -> List[str]: + async def get_datasource_priority(self) -> list[str]: if self._storage is not None: return await self._storage.get_datasource_priority() body = await self._request(method="GET", path="/datasources/priority") return list(body.get("priority", [])) - async def set_datasource_priority(self, priority: List[str]) -> None: + async def set_datasource_priority(self, priority: list[str]) -> None: if self._storage is not None: await self._storage.set_datasource_priority(list(priority)) return @@ -342,9 +341,9 @@ async def save_memory( self, *, learning: str, - linked_entities: Union[List[str], SlayerQuery, Dict[str, Any]], - id: Optional[str] = None, # noqa: A002 — public kwarg matching MCP / REST - description: Optional[str] = None, + linked_entities: list[str] | SlayerQuery | dict[str, Any], + id: str | None = None, # noqa: A002 — public kwarg matching MCP / REST + description: str | None = None, ) -> SaveMemoryResponse: """Save a memory: a learning text + linked entities (or an inline SlayerQuery to extract entities from). DEV-1428: @@ -361,7 +360,7 @@ async def save_memory( description=description, ) return response - body: Dict[str, Any] = { + body: dict[str, Any] = { "learning": learning, "linked_entities": self._coerce_linked_entities(linked_entities), } @@ -373,7 +372,7 @@ async def save_memory( return SaveMemoryResponse.model_validate(result) async def forget_memory( - self, identifier: Union[int, str] + self, identifier: int | str ) -> ForgetMemoryResponse: if self._storage is not None: return await self._memory_service().forget_memory( @@ -393,12 +392,12 @@ async def forget_memory( async def search( self, *, - entities: Optional[List[str]] = None, - query: Optional[Union[SlayerQuery, Dict[str, Any]]] = None, - question: Optional[str] = None, - datasource: Optional[str] = None, + entities: list[str] | None = None, + query: SlayerQuery | dict[str, Any] | None = None, + question: str | None = None, + datasource: str | None = None, max_results: int = 10, - cypher_filter: Optional[str] = None, + cypher_filter: str | None = None, compact: bool = True, ) -> "SearchResponse": """Up to three-channel semantic search over memories + canonical @@ -446,7 +445,7 @@ async def search( cypher_filter=cypher_filter, compact=compact, ) - body: Dict[str, Any] = { + body: dict[str, Any] = { "max_results": max_results, "compact": compact, } @@ -514,8 +513,8 @@ def query_df(self, query: QueryInput): result = self.query_sync(query=query) return pd.DataFrame(result.data) - def list_models_sync(self) -> List[str]: + def list_models_sync(self) -> list[str]: return self._request_sync(method="GET", path="/models") - def get_model_sync(self, name: str) -> Dict[str, Any]: + def get_model_sync(self, name: str) -> dict[str, Any]: return self._request_sync("GET", f"/models/{name}") diff --git a/slayer/core/enums.py b/slayer/core/enums.py index d6ee8652..707810be 100644 --- a/slayer/core/enums.py +++ b/slayer/core/enums.py @@ -2,7 +2,7 @@ import datetime # noqa: F401 (kept for downstream imports of TimeGranularity) from enum import Enum -from typing import Any, Optional +from typing import Any class StrEnum(str, Enum): @@ -27,7 +27,7 @@ class DataType(StrEnum): # from older agent input (MCP/REST/CLI), pseudo-types (count/sum/...) drop to # None so the field falls through to its default. Used by both Column and # ModelMeasure validators in slayer/core/models.py. -_LEGACY_DATATYPE_ALIASES: dict[str, Optional[str]] = { +_LEGACY_DATATYPE_ALIASES: dict[str, str | None] = { # Pre-rename canonical values. "string": "TEXT", "number": "DOUBLE", diff --git a/slayer/core/errors.py b/slayer/core/errors.py index fb0993ef..305bdb44 100644 --- a/slayer/core/errors.py +++ b/slayer/core/errors.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, List, Tuple +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from slayer.engine.schema_drift import ToDeleteEntry # noqa: F401 @@ -35,7 +35,7 @@ class AmbiguousModelError(SlayerError): priority`` CLI subcommand, etc.). """ - def __init__(self, name: str, candidates: List[str]) -> None: + def __init__(self, name: str, candidates: list[str]) -> None: self.name = name self.candidates = list(candidates) super().__init__( @@ -85,8 +85,8 @@ class SchemaDriftError(SlayerError): def __init__( self, - models: List[str], - to_delete: List[Any], + models: list[str], + to_delete: list[Any], original: BaseException, ) -> None: self.models = list(models) @@ -110,8 +110,8 @@ class ColumnCycleError(SlayerError, ValueError): compile-time cycle raise) continue to work unchanged. """ - def __init__(self, cycle: List[Tuple[str, str]]) -> None: - self.cycle: List[Tuple[str, str]] = list(cycle) + def __init__(self, cycle: list[tuple[str, str]]) -> None: + self.cycle: list[tuple[str, str]] = list(cycle) chain = " → ".join(f"{m}.{c}" for m, c in self.cycle) super().__init__(f"Circular column reference detected: {chain}") diff --git a/slayer/core/format.py b/slayer/core/format.py index 520884fd..0938a650 100644 --- a/slayer/core/format.py +++ b/slayer/core/format.py @@ -4,7 +4,6 @@ import math import numbers from enum import Enum -from typing import Optional from pydantic import BaseModel, Field, model_validator @@ -28,12 +27,12 @@ class NumberFormat(BaseModel): default=NumberFormatType.FLOAT, description="The format type for number display", ) - precision: Optional[int] = Field( + precision: int | None = Field( default=None, ge=0, description="Number of decimal places to show", ) - symbol: Optional[str] = Field( + symbol: str | None = Field( default=None, description="Currency symbol (defaults to $ for CURRENCY type, must be None otherwise)", ) @@ -51,8 +50,8 @@ def validate_symbol(self) -> "NumberFormat": def _format_with_notation( value: float, default_precision: int, - explicit_precision: Optional[int] = None, - max_precision: Optional[int] = None, + explicit_precision: int | None = None, + max_precision: int | None = None, ) -> tuple[float, str, int]: """Core formatting logic with K/M notation and dynamic precision calculation. diff --git a/slayer/core/formula.py b/slayer/core/formula.py index 4210d0ec..d41dd5a2 100644 --- a/slayer/core/formula.py +++ b/slayer/core/formula.py @@ -17,7 +17,8 @@ import re import tokenize import warnings -from typing import Any, Dict, List, Literal, Mapping, Optional, Union +from typing import Any, Literal +from collections.abc import Mapping from pydantic import BaseModel, Field @@ -150,15 +151,15 @@ class AggregatedMeasureRef(BaseModel): """ measure_name: str = Field(description="Measure name, e.g. 'revenue', 'customers.revenue', '*'") aggregation_name: str = Field(description="Aggregation name, e.g. 'sum', 'weighted_avg'") - agg_args: List[str] = Field(default_factory=list, description="Positional aggregation args") - agg_kwargs: Dict[str, str] = Field(default_factory=dict, description="Keyword aggregation args") + agg_args: list[str] = Field(default_factory=list, description="Positional aggregation args") + agg_kwargs: dict[str, str] = Field(default_factory=dict, description="Keyword aggregation args") class ArithmeticField(BaseModel): """An arithmetic expression over measures only (no transform calls inside).""" sql: str = Field(description="Preprocessed formula with placeholders for aggregated refs") - measure_names: List[str] = Field(description="Placeholder IDs or bare measure names") - agg_refs: Dict[str, AggregatedMeasureRef] = Field(default_factory=dict) + measure_names: list[str] = Field(description="Placeholder IDs or bare measure names") + agg_refs: dict[str, AggregatedMeasureRef] = Field(default_factory=dict) is_predicate: bool = Field( default=False, description="True when the top-level AST node is a comparison or boolean op " @@ -171,8 +172,8 @@ class TransformField(BaseModel): """A transform function call, possibly wrapping another transform or arithmetic.""" transform: str = Field(description="Transform name: cumsum, lag, lead, change, change_pct, rank, percent_rank, dense_rank, ntile, time_shift, first, last, consecutive_periods") inner: "FieldSpec" = Field(description="The measure or expression being transformed") - args: List[Any] = Field(default_factory=list, description="Extra transform args (offset, granularity, etc.)") - kwargs: Dict[str, Any] = Field( + args: list[Any] = Field(default_factory=list, description="Extra transform args (offset, granularity, etc.)") + kwargs: dict[str, Any] = Field( default_factory=dict, description=( "Keyword args from the call site, e.g. partition_by=[...] for the " @@ -188,9 +189,9 @@ class MixedArithmeticField(BaseModel): as a CTE step, then the arithmetic references its result. """ sql: str = Field(description="Preprocessed formula with placeholders") - measure_names: List[str] = Field(description="Placeholder IDs or bare measure names") - sub_transforms: List[tuple] = Field(description="List of (placeholder_name, TransformField)") - agg_refs: Dict[str, AggregatedMeasureRef] = Field(default_factory=dict) + measure_names: list[str] = Field(description="Placeholder IDs or bare measure names") + sub_transforms: list[tuple] = Field(description="List of (placeholder_name, TransformField)") + agg_refs: dict[str, AggregatedMeasureRef] = Field(default_factory=dict) is_predicate: bool = Field( default=False, description="True when the top-level AST node is a comparison or boolean op " @@ -200,7 +201,7 @@ class MixedArithmeticField(BaseModel): # The parsed result of a single field -FieldSpec = Union[AggregatedMeasureRef, ArithmeticField, TransformField, MixedArithmeticField] +FieldSpec = AggregatedMeasureRef | ArithmeticField | TransformField | MixedArithmeticField # Rebuild TransformField which uses a forward reference to FieldSpec TransformField.model_rebuild() @@ -243,7 +244,7 @@ def _find_balanced_close(s: str, start: int) -> int: def _rewrite_funcstyle_aggregations( formula: str, - extra_agg_names: Optional[frozenset[str]] = None, + extra_agg_names: frozenset[str] | None = None, ) -> str: """Rewrite function-style aggregation calls to colon syntax. @@ -389,7 +390,7 @@ def _split_args(s: str) -> list[str]: # --------------------------------------------------------------------------- -def _split_agg_arglist(args_str: Optional[str]) -> tuple[list[str], dict[str, str]]: +def _split_agg_arglist(args_str: str | None) -> tuple[list[str], dict[str, str]]: """Parse a colon-aggregation ``(...)`` arglist into (positional, keyword). ``args_str`` is the raw ``(...)`` capture (with parens) or ``None``. @@ -416,7 +417,7 @@ def _split_agg_arglist(args_str: Optional[str]) -> tuple[list[str], dict[str, st def _preprocess_agg_refs( formula: str, custom_agg_names: frozenset[str] = frozenset(), -) -> tuple[str, Dict[str, AggregatedMeasureRef]]: +) -> tuple[str, dict[str, AggregatedMeasureRef]]: """Replace colon-syntax aggregated measure refs with placeholder identifiers. Returns (preprocessed_formula, {placeholder: AggregatedMeasureRef}). @@ -432,7 +433,7 @@ def _preprocess_agg_refs( aggregation. Per-ref scoping would require deferring the heal past parse, which would desync the canonical filter/ORDER-BY aliases built here. """ - refs: Dict[str, AggregatedMeasureRef] = {} + refs: dict[str, AggregatedMeasureRef] = {} counter = [0] def _replace(match: re.Match) -> str: @@ -506,7 +507,7 @@ def _expand_named_measures( except (tokenize.TokenError, IndentationError, SyntaxError): return formula - def _significant(idx: int, step: int) -> Optional[tokenize.TokenInfo]: + def _significant(idx: int, step: int) -> tokenize.TokenInfo | None: skip_types = {tokenize.NEWLINE, tokenize.NL, tokenize.ENCODING, tokenize.ENDMARKER, tokenize.COMMENT, tokenize.INDENT, tokenize.DEDENT} @@ -519,7 +520,7 @@ def _significant(idx: int, step: int) -> Optional[tokenize.TokenInfo]: return t return None - replacements: List[tuple] = [] # (start_pos, end_pos, replacement_text) + replacements: list[tuple] = [] # (start_pos, end_pos, replacement_text) for i, tok in enumerate(tokens): if tok.type != tokenize.NAME or tok.string not in named_measures: continue @@ -561,8 +562,8 @@ def _significant(idx: int, step: int) -> Optional[tokenize.TokenInfo]: def parse_formula( formula: str, - extra_agg_names: Optional[frozenset[str]] = None, - named_measures: Optional[Mapping[str, str]] = None, + extra_agg_names: frozenset[str] | None = None, + named_measures: Mapping[str, str] | None = None, ) -> FieldSpec: """Parse a formula string into a FieldSpec. @@ -605,7 +606,7 @@ def parse_formula( def _parse_node( node: ast.AST, original: str, - agg_refs: Optional[Dict[str, AggregatedMeasureRef]] = None, + agg_refs: dict[str, AggregatedMeasureRef] | None = None, ) -> FieldSpec: """Recursively parse an AST node into a FieldSpec.""" if agg_refs is None: @@ -725,7 +726,7 @@ def _replace_calls_in_arith( sub_transforms: list[tuple], measure_names: list[str], counter: list[int], - agg_refs: Dict[str, AggregatedMeasureRef], + agg_refs: dict[str, AggregatedMeasureRef], original: str, ) -> ast.AST: """Walk the AST, replacing transform Call nodes with Name placeholders.""" @@ -793,7 +794,7 @@ def _replace_calls_in_arith( def _parse_mixed_arithmetic( node: ast.AST, original: str, - agg_refs: Optional[Dict[str, AggregatedMeasureRef]] = None, + agg_refs: dict[str, AggregatedMeasureRef] | None = None, ) -> MixedArithmeticField: """Parse arithmetic that contains transform calls. @@ -839,7 +840,7 @@ def _parse_literal(node: ast.AST, original: str) -> Any: # Per-transform kwarg whitelist. Empty set means the transform takes no kwargs. -_ALLOWED_TRANSFORM_KWARGS: Dict[str, frozenset] = { +_ALLOWED_TRANSFORM_KWARGS: dict[str, frozenset] = { "rank": frozenset({"partition_by"}), "percent_rank": frozenset({"partition_by"}), "dense_rank": frozenset({"partition_by"}), @@ -863,8 +864,8 @@ def _parse_dotted_name(node: ast.AST, original: str) -> str: def _parse_transform_kwargs( # NOSONAR S3776 — straight-line whitelist + per-kwarg validation; splitting into helpers would force threading transform/original through every call just to preserve the error-message context - transform: str, keywords: List[ast.keyword], original: str -) -> Dict[str, Any]: + transform: str, keywords: list[ast.keyword], original: str +) -> dict[str, Any]: """Parse and validate a transform's keyword arguments. Each transform has a fixed kwarg whitelist (see ``_ALLOWED_TRANSFORM_KWARGS``). @@ -873,7 +874,7 @@ def _parse_transform_kwargs( # NOSONAR S3776 — straight-line whitelist + per- path, or list of those. """ allowed = _ALLOWED_TRANSFORM_KWARGS.get(transform, frozenset()) - parsed: Dict[str, Any] = {} + parsed: dict[str, Any] = {} for kw in keywords: if kw.arg is None: @@ -941,10 +942,10 @@ class ParsedFilter(BaseModel): as-is (they get qualified with the model name during SQL generation). """ sql: str = Field(description="SQL WHERE condition, e.g. \"status = 'completed'\"") - columns: List[str] = Field(description="Column names referenced in the filter") + columns: list[str] = Field(description="Column names referenced in the filter") is_having: bool = Field(default=False, description="True if this is a HAVING filter (aggregate condition)") is_post_filter: bool = Field(default=False, description="True if this references a computed column (transform/expression)") - synthesized_aliases: List[str] = Field( + synthesized_aliases: list[str] = Field( default_factory=list, description=( "Canonical aggregation aliases this filter introduced from " @@ -954,7 +955,7 @@ class ParsedFilter(BaseModel): "regex that would let typos like ``made_up_sum`` through." ), ) - agg_refs: List["AggregatedMeasureRef"] = Field( + agg_refs: list["AggregatedMeasureRef"] = Field( default_factory=list, description=( "Aggregated measure references extracted from colon syntax in " @@ -1031,7 +1032,7 @@ def _preprocess_concat(formula: str) -> str: parts = _STRING_LITERAL_RE.split(formula) literals = _STRING_LITERAL_RE.findall(formula) rewritten_parts = [p.replace("||", "<<") for p in parts] - out: List[str] = [] + out: list[str] = [] for i, p in enumerate(rewritten_parts): out.append(p) if i < len(literals): @@ -1068,8 +1069,8 @@ def _preprocess_sql_operators(formula: str) -> str: def parse_filter( formula: str, - extra_agg_names: Optional[frozenset[str]] = None, - named_measures: Optional[Mapping[str, str]] = None, + extra_agg_names: frozenset[str] | None = None, + named_measures: Mapping[str, str] | None = None, ) -> ParsedFilter: """Parse a Mode B (DSL) filter formula into a ParsedFilter. @@ -1167,7 +1168,7 @@ def parse_filter( ) -_BINOP_OP_MAP: Dict[type, str] = { +_BINOP_OP_MAP: dict[type, str] = { ast.Add: "+", ast.Sub: "-", ast.Mult: "*", ast.Div: "/", ast.Mod: "%", ast.Pow: "**", } @@ -1185,7 +1186,7 @@ def parse_filter( # That adds at most one harmless paren on mixed-precedence Pow # expressions and guarantees ``(a ** b) ** c`` doesn't silently # re-associate to ``a ** (b ** c)``. -_BINOP_PRECEDENCE: Dict[type, int] = { +_BINOP_PRECEDENCE: dict[type, int] = { ast.Mult: 2, ast.Div: 2, ast.Mod: 2, ast.Add: 1, ast.Sub: 1, } @@ -1295,7 +1296,7 @@ def _seq_to_sql(node, recur) -> str: return f"({', '.join(elts)})" -def _flatten_lshift_chain(node: ast.AST, recur) -> List[str]: +def _flatten_lshift_chain(node: ast.AST, recur) -> list[str]: """Flatten a chain of ``ast.BinOp(LShift)`` nodes into a flat list of SQL strings. Used by ``_binop_to_sql`` to fold a ``a || b || c`` chain (which arrives as left-associative LShift after `_preprocess_concat`) @@ -1332,7 +1333,7 @@ def _binop_to_sql(node: ast.BinOp, original: str, recur) -> str: def _emit_binop_operand( child: ast.AST, - parent_prec: Optional[int], + parent_prec: int | None, *, is_right: bool, recur, @@ -1488,7 +1489,7 @@ def _get_string_arg(node: ast.AST, original: str) -> str: raise ValueError(f"Expected a string argument in filter: {original!r}") -def _collect_names(node: ast.AST) -> List[str]: +def _collect_names(node: ast.AST) -> list[str]: """Collect all Name and dotted Attribute references from an AST subtree.""" names = [] # Collect dotted names (model.measure) first to avoid also collecting the bare Name part diff --git a/slayer/core/models.py b/slayer/core/models.py index ac5510f1..34615919 100644 --- a/slayer/core/models.py +++ b/slayer/core/models.py @@ -3,7 +3,7 @@ import logging import os import re -from typing import Annotated, Any, Dict, List, Optional +from typing import Annotated, Any, Optional from pydantic import BaseModel, BeforeValidator, Field, field_validator, model_validator from sqlalchemy.engine import URL as _SA_URL @@ -161,19 +161,19 @@ class Column(BaseModel): Replaces v1 ``Dimension`` and ``Measure`` (which were merged in v2). """ name: str - sql: Optional[str] = None + sql: str | None = None type: DataType = DataType.TEXT primary_key: bool = False - description: Optional[str] = None - label: Optional[str] = None + description: str | None = None + label: str | None = None hidden: bool = False - format: Optional[NumberFormat] = None - allowed_aggregations: Optional[List[str]] = None - filter: Optional[str] = None # Applied inside CASE WHEN at aggregation time only - meta: Optional[Dict[str, Any]] = None - sampled: Optional[str] = None # DEV-1375: cached sample-value snapshot - sampled_values: Optional[List[str]] = None # DEV-1480: structured top-N - distinct_count: Optional[int] = None # DEV-1480: true cardinality at profile time + format: NumberFormat | None = None + allowed_aggregations: list[str] | None = None + filter: str | None = None # Applied inside CASE WHEN at aggregation time only + meta: dict[str, Any] | None = None + sampled: str | None = None # DEV-1375: cached sample-value snapshot + sampled_values: list[str] | None = None # DEV-1480: structured top-N + distinct_count: int | None = None # DEV-1480: true cardinality at profile time @model_validator(mode="before") @classmethod @@ -196,14 +196,14 @@ def _validate_name(cls, v: str) -> str: @field_validator("sql") @classmethod - def _fix_multidot_sql(cls, v: Optional[str]) -> Optional[str]: + def _fix_multidot_sql(cls, v: str | None) -> str | None: if v is not None: v = _fix_multidot_sql(v, context="Column sql") return v @field_validator("filter") @classmethod - def _fix_multidot_filter(cls, v: Optional[str]) -> Optional[str]: + def _fix_multidot_filter(cls, v: str | None) -> str | None: if v is not None: v = _fix_multidot_sql(v, context="Column filter") # DEV-1369: Column.filter is SQL-mode — validate at construction @@ -228,11 +228,11 @@ class ModelMeasure(BaseModel): contexts; the difference is scope. """ formula: str - name: Optional[str] = None - label: Optional[str] = None - description: Optional[str] = None - type: Optional[DataType] = None - meta: Optional[Dict[str, Any]] = None + name: str | None = None + label: str | None = None + description: str | None = None + type: DataType | None = None + meta: dict[str, Any] | None = None @model_validator(mode="before") @classmethod @@ -250,7 +250,7 @@ def _coerce_legacy_type(cls, data: Any) -> Any: @field_validator("name") @classmethod - def _validate_name(cls, v: Optional[str]) -> Optional[str]: + def _validate_name(cls, v: str | None) -> str | None: if v is not None and not _NAME_PATTERN.match(v): raise ValueError( f"Invalid name '{v}': must contain only letters, digits, " @@ -260,7 +260,7 @@ def _validate_name(cls, v: Optional[str]) -> Optional[str]: @field_validator("name") @classmethod - def _reject_transform_shadowing(cls, v: Optional[str]) -> Optional[str]: + def _reject_transform_shadowing(cls, v: str | None) -> str | None: """A saved measure named after a built-in transform (``cumsum`` etc.) would shadow the transform when written as ``cumsum(...)`` in another formula. Reject these names at construction time. @@ -308,10 +308,10 @@ class Aggregation(BaseModel): For fully custom aggregations, ``formula`` is required. """ name: str - formula: Optional[str] = None # SQL template; None = use built-in formula - params: List[AggregationParam] = Field(default_factory=list) - description: Optional[str] = None - meta: Optional[Dict[str, Any]] = None + formula: str | None = None # SQL template; None = use built-in formula + params: list[AggregationParam] = Field(default_factory=list) + description: str | None = None + meta: dict[str, Any] | None = None @field_validator("name") @classmethod @@ -408,7 +408,7 @@ class SourceModelOrigin(BaseModel): ``exclude=True`` so accidental save paths drop it cleanly. """ name: str - data_source: Optional[str] = None + data_source: str | None = None parent: Optional["SourceModelOrigin"] = None agg_column_names: frozenset[str] = Field(default_factory=frozenset) @@ -416,12 +416,12 @@ class SourceModelOrigin(BaseModel): class ModelJoin(BaseModel): """A join relationship to another model.""" target_model: str # Name of the joined model - join_pairs: List[List[str]] = Field(...) # [["source_dim", "target_dim"], ...] + join_pairs: list[list[str]] = Field(...) # [["source_dim", "target_dim"], ...] join_type: JoinType = JoinType.LEFT # LEFT (default) or INNER @field_validator("join_pairs") @classmethod - def _validate_join_pairs(cls, v: List[List[str]]) -> List[List[str]]: + def _validate_join_pairs(cls, v: list[list[str]]) -> list[list[str]]: if not v: raise ValueError("join_pairs must be non-empty") for i, pair in enumerate(v): @@ -435,17 +435,17 @@ def _validate_join_pairs(cls, v: List[List[str]]) -> List[List[str]]: class SlayerModel(BaseModel): version: int = 7 name: str - sql_table: Optional[str] = None - sql: Optional[str] = None + sql_table: str | None = None + sql: str | None = None source_queries: Annotated[ - Optional[List], BeforeValidator(_coerce_source_queries) + list | None, BeforeValidator(_coerce_source_queries) ] = None # List of SlayerQuery — query-backed source mode - query_variables: Dict[str, Any] = Field(default_factory=dict) - backing_query_sql: Optional[str] = None + query_variables: dict[str, Any] = Field(default_factory=dict) + backing_query_sql: str | None = None data_source: str = "" - columns: List[Column] = Field(default_factory=list) - measures: List[ModelMeasure] = Field(default_factory=list) - aggregations: List[Aggregation] = Field(default_factory=list) + columns: list[Column] = Field(default_factory=list) + measures: list[ModelMeasure] = Field(default_factory=list) + aggregations: list[Aggregation] = Field(default_factory=list) @model_validator(mode="before") @classmethod @@ -497,20 +497,20 @@ def _require_data_source_unless_query_backed(self) -> "SlayerModel": f"model belongs to." ) return self - joins: List[ModelJoin] = Field(default_factory=list) - filters: List[str] = Field(default_factory=list) # Model-level filters (always applied) - default_time_dimension: Optional[str] = None - description: Optional[str] = None + joins: list[ModelJoin] = Field(default_factory=list) + filters: list[str] = Field(default_factory=list) # Model-level filters (always applied) + default_time_dimension: str | None = None + description: str | None = None hidden: bool = False - meta: Optional[Dict[str, Any]] = None + meta: dict[str, Any] | None = None # DEV-1449: in-memory breadcrumb for virtual stage models produced by # ``_query_as_model``. ``exclude=True`` keeps it out of YAML/SQLite # roundtrips; virtual stage models are not persisted in the first place. - source_model_origin: Optional[SourceModelOrigin] = Field(default=None, exclude=True) + source_model_origin: SourceModelOrigin | None = Field(default=None, exclude=True) @field_validator("filters") @classmethod - def _fix_multidot_filters(cls, v: List[str]) -> List[str]: + def _fix_multidot_filters(cls, v: list[str]) -> list[str]: """Auto-convert multi-dot column references in model filters and validate each entry as a SQL-mode predicate (DEV-1369). @@ -685,7 +685,7 @@ def _validate_source_query_stages(self) -> "SlayerModel": f"in source_queries must have a 'name'." ) seen: set = set() - dupes: List[str] = [] + dupes: list[str] = [] for stage in stages: n = getattr(stage, "name", None) if not n: @@ -700,19 +700,19 @@ def _validate_source_query_stages(self) -> "SlayerModel": ) return self - def get_column(self, name: str) -> Optional[Column]: + def get_column(self, name: str) -> Column | None: for c in self.columns: if c.name == name: return c return None - def get_measure(self, name: str) -> Optional[ModelMeasure]: + def get_measure(self, name: str) -> ModelMeasure | None: for m in self.measures: if m.name == name: return m return None - def get_aggregation(self, name: str) -> Optional[Aggregation]: + def get_aggregation(self, name: str) -> Aggregation | None: for a in self.aggregations: if a.name == name: return a @@ -722,15 +722,15 @@ def get_aggregation(self, name: str) -> Optional[Aggregation]: class DatasourceConfig(BaseModel): version: int = 2 name: str - type: Optional[str] = None - host: Optional[str] = None - port: Optional[int] = None - database: Optional[str] = None - username: Optional[str] = None - password: Optional[str] = None - connection_string: Optional[str] = None - schema_name: Optional[str] = None - description: Optional[str] = None + type: str | None = None + host: str | None = None + port: int | None = None + database: str | None = None + username: str | None = None + password: str | None = None + connection_string: str | None = None + schema_name: str | None = None + description: str | None = None # Snowflake-specific (v2, DEV-1551). Other dialects ignore them. # ``connection_name`` is the primary auth path — when set, all credentials # come from ``~/.snowflake/connections.toml`` via @@ -738,16 +738,16 @@ class DatasourceConfig(BaseModel): # (host as account, username, password, database, schema_name) is the # secondary path; ``warehouse`` and ``role`` populate the URL's query # string for that path. See docs/configuration/datasources.md#snowflake. - connection_name: Optional[str] = None - warehouse: Optional[str] = None - role: Optional[str] = None + connection_name: str | None = None + warehouse: str | None = None + role: str | None = None # BigQuery-specific. Other dialects ignore it. The contents of a Google # service-account key file as a JSON string. When set, the BigQuery dialect # constructs the SQLAlchemy engine with ``credentials_info=json.loads(...)`` # so the connection authenticates against that service account directly. # When unset, BigQuery falls back to Application Default Credentials # (``GOOGLE_APPLICATION_CREDENTIALS`` env var or attached compute identity). - credentials_json: Optional[str] = Field(default=None, repr=False) + credentials_json: str | None = Field(default=None, repr=False) @model_validator(mode="before") @classmethod diff --git a/slayer/core/query.py b/slayer/core/query.py index 5dae7054..da3c8fe3 100644 --- a/slayer/core/query.py +++ b/slayer/core/query.py @@ -9,7 +9,7 @@ import datetime import logging import re -from typing import Annotated, Any, Dict, List, Optional +from typing import Annotated, Any from pydantic import BaseModel, BeforeValidator, ConfigDict, field_validator, model_validator @@ -38,7 +38,7 @@ def _validate_query_filter_string(formula: str) -> None: raise ValueError(f"Filter '{formula}' {WINDOW_IN_FILTER_ERROR}") -def substitute_variables(filter_str: str, variables: Dict[str, Any]) -> str: +def substitute_variables(filter_str: str, variables: dict[str, Any]) -> str: """Substitute {variable} placeholders in a filter string. - {var_name} is replaced with the variable's value (str or number, inserted as-is). @@ -110,8 +110,8 @@ class ColumnRef(BaseModel): on the query's model. """ name: str - model: Optional[str] = None - label: Optional[str] = None + model: str | None = None + label: str | None = None @model_validator(mode="after") def _parse_dotted_name(self) -> "ColumnRef": @@ -218,8 +218,8 @@ def _is_direction(value: Any) -> bool: class TimeDimension(BaseModel): dimension: Annotated[ColumnRef, BeforeValidator(_coerce_column_ref)] granularity: TimeGranularity - date_range: Optional[List[str]] = None - label: Optional[str] = None + date_range: list[str] | None = None + label: str | None = None class OrderItem(BaseModel): @@ -230,7 +230,7 @@ class OrderItem(BaseModel): column: Annotated[ColumnRef, BeforeValidator(_coerce_order_column)] direction: str = "asc" - raw_formula: Optional[str] = None + raw_formula: str | None = None @model_validator(mode="before") @classmethod @@ -340,12 +340,12 @@ class ModelExtension(BaseModel): stored model. """ source_name: str # Model/query to extend - columns: Optional[List] = None # Extra Column objects - measures: Optional[List[ModelMeasure]] = None # Extra ModelMeasure formulas - joins: Optional[List] = None # Extra ModelJoin objects + columns: list | None = None # Extra Column objects + measures: list[ModelMeasure] | None = None # Extra ModelMeasure formulas + joins: list | None = None # Extra ModelJoin objects -def _get_source_model_name(source_model: object) -> Optional[str]: +def _get_source_model_name(source_model: object) -> str | None: """Extract the model name from any source_model type. Works before model resolution — handles str, dict, ModelExtension, @@ -399,9 +399,9 @@ class SlayerQuery(BaseModel): model_config = ConfigDict(extra="forbid") version: int = 3 - name: Optional[str] = None # For referencing this query from other queries in a list + name: str | None = None # For referencing this query from other queries in a list source_model: object # str (model name), SlayerModel (inline), or ModelExtension - measures: Annotated[Optional[List[ModelMeasure]], BeforeValidator(_coerce_measures)] = None + measures: Annotated[list[ModelMeasure] | None, BeforeValidator(_coerce_measures)] = None @model_validator(mode="before") @classmethod @@ -410,7 +410,7 @@ def _apply_schema_migrations(cls, data: Any) -> Any: @field_validator("name") @classmethod - def _validate_query_name(cls, v: Optional[str]) -> Optional[str]: + def _validate_query_name(cls, v: str | None) -> str | None: # Share the same rejection rules as SlayerModel.name — # SlayerQuery names occupy the same naming space when persisted # as query-backed models. Rejects ``__`` (join-path alias @@ -420,14 +420,14 @@ def _validate_query_name(cls, v: Optional[str]) -> Optional[str]: return v from slayer.core.models import _validate_model_name return _validate_model_name(v, "Query") - dimensions: Annotated[Optional[List[ColumnRef]], BeforeValidator(_coerce_dimensions)] = None - time_dimensions: Optional[List[TimeDimension]] = None - main_time_dimension: Optional[str] = None # Explicit time dimension for transforms (overrides auto-detection) - filters: Optional[List[str]] = None - variables: Optional[Dict[str, Any]] = None # Variable values for filter substitution - order: Annotated[Optional[List[OrderItem]], BeforeValidator(_coerce_order)] = None - limit: Optional[int] = None - offset: Optional[int] = None + dimensions: Annotated[list[ColumnRef] | None, BeforeValidator(_coerce_dimensions)] = None + time_dimensions: list[TimeDimension] | None = None + main_time_dimension: str | None = None # Explicit time dimension for transforms (overrides auto-detection) + filters: list[str] | None = None + variables: dict[str, Any] | None = None # Variable values for filter substitution + order: Annotated[list[OrderItem] | None, BeforeValidator(_coerce_order)] = None + limit: int | None = None + offset: int | None = None whole_periods_only: bool = False # DEV-1543: opt out of the auto "distinct dimension tuples" GROUP BY # for dim-only queries. Default ``True`` preserves the Cube.js-style @@ -534,7 +534,7 @@ def strip_source_model_prefix(self) -> "SlayerQuery": if model_name is None: return self - updates: Dict[str, Any] = {} + updates: dict[str, Any] = {} pattern = re.compile(r"\b" + re.escape(model_name) + r"\.") # Dimensions diff --git a/slayer/core/refs.py b/slayer/core/refs.py index 01aa8fb8..beae3706 100644 --- a/slayer/core/refs.py +++ b/slayer/core/refs.py @@ -13,7 +13,6 @@ from __future__ import annotations import re -from typing import List, Optional, Tuple # --------------------------------------------------------------------------- # Identifier shapes @@ -54,8 +53,8 @@ def agg_signature_suffix( - agg_args: Optional[List[str]], - agg_kwargs: Optional[dict], + agg_args: list[str] | None, + agg_kwargs: dict | None, ) -> str: """Build a deterministic identifier suffix from aggregation args/kwargs. @@ -70,7 +69,7 @@ def agg_signature_suffix( kwargs = agg_kwargs or {} if not args and not kwargs: return "" - parts: List[str] = [] + parts: list[str] = [] for a in args: sanitized = _NON_IDENT_RE.sub("_", str(a)).strip("_") if sanitized: @@ -88,8 +87,8 @@ def agg_signature_suffix( def canonical_agg_name( measure_name: str, aggregation_name: str, - agg_args: Optional[List[str]] = None, - agg_kwargs: Optional[dict] = None, + agg_args: list[str] | None = None, + agg_kwargs: dict | None = None, ) -> str: """Canonical hidden-column name for an aggregated measure ref. @@ -103,7 +102,7 @@ def canonical_agg_name( return f"{measure_name}_{aggregation_name}{suffix}" -def strip_agg_suffix(raw: str) -> Tuple[str, Optional[str]]: +def strip_agg_suffix(raw: str) -> tuple[str, str | None]: """Return ``(prefix, agg_name)`` after stripping a trailing ``:agg`` or ``:agg(...)``. diff --git a/slayer/dbt/converter.py b/slayer/dbt/converter.py index 0f3bb545..b2449088 100644 --- a/slayer/dbt/converter.py +++ b/slayer/dbt/converter.py @@ -16,7 +16,6 @@ import logging import re from collections import defaultdict -from typing import Dict, List, Optional, Tuple import sqlalchemy as sa from pydantic import BaseModel, Field @@ -42,7 +41,7 @@ logger = logging.getLogger(__name__) # Map dbt aggregation names to SLayer aggregation names -_AGG_MAP: Dict[str, str] = { +_AGG_MAP: dict[str, str] = { "sum": "sum", "average": "avg", "avg": "avg", @@ -68,16 +67,16 @@ class DbtConversionError(Exception): class ConversionWarning(BaseModel): """A warning or info message from the conversion process.""" - model_name: Optional[str] = None - metric_name: Optional[str] = None + model_name: str | None = None + metric_name: str | None = None message: str class ConversionResult(BaseModel): """Result of converting a DbtProject to SLayer representations.""" - models: List[SlayerModel] = Field(default_factory=list) - unconverted_metrics: List[ConversionWarning] = Field(default_factory=list) - warnings: List[ConversionWarning] = Field(default_factory=list) + models: list[SlayerModel] = Field(default_factory=list) + unconverted_metrics: list[ConversionWarning] = Field(default_factory=list) + warnings: list[ConversionWarning] = Field(default_factory=list) def _map_agg(dbt_agg: str) -> str: @@ -113,12 +112,12 @@ def _convert_dimension(dim: DbtDimension) -> Column: def _convert_measures( - dbt_measures: List[DbtMeasure], + dbt_measures: list[DbtMeasure], *, sm_name: str, existing_column_names: set, - unconverted: List[ConversionWarning], -) -> Tuple[List[Column], List[ModelMeasure]]: + unconverted: list[ConversionWarning], +) -> tuple[list[Column], list[ModelMeasure]]: """Convert dbt measures into a (Columns, ModelMeasures) pair. Each unique measure expression yields a single ``Column`` whose name is @@ -137,7 +136,7 @@ def _convert_measures( ``cumsum``) is rejected by the Pydantic validator; the dbt measure is routed to ``unconverted`` and skipped. """ - groups: Dict[str, List[DbtMeasure]] = defaultdict(list) + groups: dict[str, list[DbtMeasure]] = defaultdict(list) for m in dbt_measures: key = m.expr or m.name groups[key].append(m) @@ -145,8 +144,8 @@ def _convert_measures( # All dbt-measure names, taken to be the eventual ``ModelMeasure`` names. measure_names = {m.name for m in dbt_measures} - columns: List[Column] = [] - measures: List[ModelMeasure] = [] + columns: list[Column] = [] + measures: list[ModelMeasure] = [] used_column_names = set(existing_column_names) for expr_key, group in groups.items(): @@ -200,7 +199,7 @@ def __init__( self, project: DbtProject, data_source: str, - sa_engine: Optional[sa.Engine] = None, + sa_engine: sa.Engine | None = None, include_hidden_models: bool = False, ) -> None: self.project = project @@ -208,15 +207,15 @@ def __init__( self.sa_engine = sa_engine self.include_hidden_models = include_hidden_models self.entity_registry = EntityRegistry() - self._warnings: List[ConversionWarning] = [] - self._unconverted: List[ConversionWarning] = [] + self._warnings: list[ConversionWarning] = [] + self._unconverted: list[ConversionWarning] = [] # {model_name: SlayerModel} for metric resolution - self._models_by_name: Dict[str, SlayerModel] = {} + self._models_by_name: dict[str, SlayerModel] = {} # {model_name: DbtSemanticModel} for looking up entities - self._dbt_models_by_name: Dict[str, DbtSemanticModel] = {} + self._dbt_models_by_name: dict[str, DbtSemanticModel] = {} # {regular_model_name: raw_code} — used to inline SQL into semantic # models whose underlying dbt model is a query rather than a table. - self._regular_models_sql: Dict[str, str] = { + self._regular_models_sql: dict[str, str] = { rm.name: rm.raw_code for rm in project.regular_models if rm.raw_code @@ -229,7 +228,7 @@ def convert(self) -> ConversionResult: for sm in self.project.semantic_models: self._dbt_models_by_name[sm.name] = sm - models: List[SlayerModel] = [] + models: list[SlayerModel] = [] for sm in self.project.semantic_models: model = self._convert_semantic_model(sm) models.append(model) @@ -270,7 +269,7 @@ def _mirror_inner_joins(self) -> None: join_type=JoinType.INNER, )) - def _convert_regular_models(self, existing_names: set) -> List[SlayerModel]: + def _convert_regular_models(self, existing_names: set) -> list[SlayerModel]: """Convert orphan dbt models (not wrapped by semantic_models) to hidden SLayer models.""" if self.sa_engine is None: self._warnings.append(ConversionWarning( @@ -283,7 +282,7 @@ def _convert_regular_models(self, existing_names: set) -> List[SlayerModel]: engine = self.sa_engine inspector = sa.inspect(engine) - results: List[SlayerModel] = [] + results: list[SlayerModel] = [] for rm in self.project.regular_models: if rm.name in existing_names: continue @@ -298,7 +297,7 @@ def _convert_regular_model( rm: DbtRegularModel, sa_engine: sa.Engine, inspector: sa.engine.Inspector, - ) -> Optional[SlayerModel]: + ) -> SlayerModel | None: """Introspect a regular dbt model and wrap it as a hidden SlayerModel.""" table_name = rm.alias or rm.name try: @@ -353,8 +352,8 @@ def _convert_semantic_model(self, sm: DbtSemanticModel) -> SlayerModel: ref_name = sm.model or sm.name - sql_source: Optional[str] = None - sql_table: Optional[str] = None + sql_source: str | None = None + sql_table: str | None = None if ref_name in self._regular_models_sql: resolved, warnings = resolve_refs( self._regular_models_sql[ref_name], @@ -373,7 +372,7 @@ def _convert_semantic_model(self, sm: DbtSemanticModel) -> SlayerModel: if sm.defaults and sm.defaults.agg_time_dimension: default_time_dim = sm.defaults.agg_time_dimension - cols: List[Column] = [_convert_dimension(d) for d in sm.dimensions] + cols: list[Column] = [_convert_dimension(d) for d in sm.dimensions] # Add primary key column for primary/unique entities. entity_col_names = {c.name for c in cols} @@ -743,7 +742,7 @@ def _convert_cumulative_metric(self, metric: DbtMetric) -> None: # ── Resolution helpers ──────────────────────────────────────────── - def _find_measure_model(self, measure_name: str) -> Optional[DbtSemanticModel]: + def _find_measure_model(self, measure_name: str) -> DbtSemanticModel | None: """Find which dbt semantic model contains a given measure.""" for sm in self.project.semantic_models: for m in sm.measures: @@ -751,7 +750,7 @@ def _find_measure_model(self, measure_name: str) -> Optional[DbtSemanticModel]: return sm return None - def _find_metric_source_model(self, metric: DbtMetric) -> Optional[str]: + def _find_metric_source_model(self, metric: DbtMetric) -> str | None: """Determine the source model for a metric. Walks ``measure``, ``metrics``, and ``numerator``/``denominator`` and @@ -765,7 +764,7 @@ def _find_metric_source_model(self, metric: DbtMetric) -> Optional[str]: sources = self._collect_metric_sources_from_params(metric.type_params) return next(iter(sources)) if len(sources) == 1 else None - def _collect_metric_sources(self, metric_name: str, _seen: Optional[set] = None) -> set: + def _collect_metric_sources(self, metric_name: str, _seen: set | None = None) -> set: """Collect every distinct semantic-model name a metric ultimately resolves to. Recurses through derived (``metrics``) and ratio @@ -789,7 +788,7 @@ def _collect_metric_sources(self, metric_name: str, _seen: Optional[set] = None) return {sm.name} if sm else set() def _collect_metric_sources_from_params( - self, type_params: DbtMetricTypeParams, *, seen: Optional[set] = None + self, type_params: DbtMetricTypeParams, *, seen: set | None = None ) -> set: """Shared shape-walker used by both entry points above.""" sources: set = set() @@ -806,7 +805,7 @@ def _collect_metric_sources_from_params( sources |= self._collect_metric_sources(side.name, _seen=seen) return sources - def _resolve_metric_to_name(self, metric_name: str) -> Optional[str]: + def _resolve_metric_to_name(self, metric_name: str) -> str | None: """Resolve a metric name to a formula reference. Returns the bare ``ModelMeasure`` name when the metric was lowered @@ -833,7 +832,7 @@ def _resolve_metric_to_name(self, metric_name: str) -> Optional[str]: return metric_name return self._resolve_measure_to_name(metric_name) - def _resolve_measure_to_name(self, measure_name: str) -> Optional[str]: + def _resolve_measure_to_name(self, measure_name: str) -> str | None: """Resolve a dbt measure name to a formula reference. After ``_convert_measures`` has run, the dbt measure name is the diff --git a/slayer/dbt/entities.py b/slayer/dbt/entities.py index b3329133..e794c355 100644 --- a/slayer/dbt/entities.py +++ b/slayer/dbt/entities.py @@ -6,7 +6,6 @@ """ import logging -from typing import Dict, List, Optional, Tuple from slayer.core.enums import JoinType from slayer.core.models import ModelJoin @@ -20,9 +19,9 @@ class EntityRegistry: def __init__(self) -> None: # {entity_name: [(model_name, expr), ...]} - self._primaries: Dict[str, List[Tuple[str, str]]] = {} + self._primaries: dict[str, list[tuple[str, str]]] = {} - def build(self, models: List[DbtSemanticModel]) -> None: + def build(self, models: list[DbtSemanticModel]) -> None: """First pass: register all primary and unique entities.""" for model in models: # Check primary_entity shorthand @@ -61,7 +60,7 @@ def _register(self, entity_name: str, model_name: str, expr: str) -> None: ) self._primaries[entity_name].append((model_name, expr)) - def get_primary_model(self, entity_name: str) -> Optional[Tuple[str, str]]: + def get_primary_model(self, entity_name: str) -> tuple[str, str] | None: """Look up which model owns this entity as primary. Returns (model_name, expr) or None. When multiple models share the @@ -73,13 +72,13 @@ def get_primary_model(self, entity_name: str) -> Optional[Tuple[str, str]]: return None return min(entries, key=lambda e: e[0]) - def resolve_joins_for_model(self, model: DbtSemanticModel) -> List[ModelJoin]: + def resolve_joins_for_model(self, model: DbtSemanticModel) -> list[ModelJoin]: """For each foreign entity in the model, generate a ModelJoin to the primary model. Returns a list of ModelJoin objects. Skips entities whose primary model is the same as the current model (self-joins are not useful). """ - joins: List[ModelJoin] = [] + joins: list[ModelJoin] = [] # Dedupe by full join signature (target + FK columns) so distinct FKs # to the same target — e.g. buyer_id -> users.id AND seller_id -> users.id — # each get their own ModelJoin instead of silently collapsing. @@ -135,14 +134,14 @@ def resolve_joins_for_model(self, model: DbtSemanticModel) -> List[ModelJoin]: return joins - def resolve_entity_to_model(self, entity_name: str) -> Optional[str]: + def resolve_entity_to_model(self, entity_name: str) -> str | None: """Given an entity name, return the first model that owns it as primary.""" entry = self.get_primary_model(entity_name) if entry is None: return None return entry[0] - def get_entity_expr(self, entity_name: str) -> Optional[str]: + def get_entity_expr(self, entity_name: str) -> str | None: """Get the SQL expression for an entity's primary key column.""" entry = self.get_primary_model(entity_name) if entry is None: diff --git a/slayer/dbt/filters.py b/slayer/dbt/filters.py index 10e0e7b1..e69f4a52 100644 --- a/slayer/dbt/filters.py +++ b/slayer/dbt/filters.py @@ -12,7 +12,6 @@ import logging import re -from typing import Dict, Optional from slayer.dbt.entities import EntityRegistry from slayer.dbt.models import DbtSemanticModel @@ -36,8 +35,8 @@ def convert_dbt_filter( filter_str: str, source_model_name: str, entity_registry: EntityRegistry, - model_entity_names: Optional[Dict[str, str]] = None, - all_semantic_models: Optional[Dict[str, DbtSemanticModel]] = None, + model_entity_names: dict[str, str] | None = None, + all_semantic_models: dict[str, DbtSemanticModel] | None = None, ) -> str: """Convert a dbt Jinja filter string to a SLayer filter string. diff --git a/slayer/dbt/manifest.py b/slayer/dbt/manifest.py index d7331ee6..73ee93ef 100644 --- a/slayer/dbt/manifest.py +++ b/slayer/dbt/manifest.py @@ -14,7 +14,7 @@ import importlib.util import logging import os -from typing import Any, Dict, List, Optional, Set +from typing import Any from slayer.dbt.models import DbtColumnMeta, DbtRegularModel @@ -34,7 +34,7 @@ def _manifest_path(project_path: str) -> str: return os.path.join(project_path, "target", "manifest.json") -def _load_manifest_file(path: str) -> Optional[dict]: +def _load_manifest_file(path: str) -> dict | None: try: with open(path, encoding="utf-8") as f: return json.load(f) @@ -67,7 +67,7 @@ def _run_dbt_parse(project_path: str) -> bool: return bool(success) -def load_or_generate_manifest(project_path: str) -> Optional[dict]: +def load_or_generate_manifest(project_path: str) -> dict | None: """Return the dbt manifest dict for a project, or None. Resolution order: @@ -95,9 +95,9 @@ def load_or_generate_manifest(project_path: str) -> Optional[dict]: return _load_manifest_file(path) -def _semantic_model_referenced_nodes(manifest: dict) -> Set[str]: +def _semantic_model_referenced_nodes(manifest: dict) -> set[str]: """Collect every dbt node key referenced by any semantic_model.""" - referenced: Set[str] = set() + referenced: set[str] = set() semantic_models = manifest.get("semantic_models") or {} for sm in semantic_models.values(): # Prefer node_relation when present, fall back to depends_on.nodes @@ -108,11 +108,11 @@ def _semantic_model_referenced_nodes(manifest: dict) -> Set[str]: return referenced -def find_orphan_model_nodes(manifest: dict) -> List[dict]: +def find_orphan_model_nodes(manifest: dict) -> list[dict]: """Return manifest nodes for regular models not wrapped by any semantic_model.""" referenced = _semantic_model_referenced_nodes(manifest) nodes = manifest.get("nodes") or {} - orphans: List[dict] = [] + orphans: list[dict] = [] for node_key, node in nodes.items(): if node.get("resource_type") != "model": continue @@ -122,7 +122,7 @@ def find_orphan_model_nodes(manifest: dict) -> List[dict]: return orphans -def _column_from_manifest(raw: Dict[str, Any]) -> DbtColumnMeta: +def _column_from_manifest(raw: dict[str, Any]) -> DbtColumnMeta: return DbtColumnMeta( name=raw.get("name", ""), description=raw.get("description") or None, @@ -131,7 +131,7 @@ def _column_from_manifest(raw: Dict[str, Any]) -> DbtColumnMeta: ) -def _regular_model_from_node(node: Dict[str, Any]) -> DbtRegularModel: +def _regular_model_from_node(node: dict[str, Any]) -> DbtRegularModel: columns_raw = node.get("columns") or {} columns = [_column_from_manifest(c) for c in columns_raw.values() if c.get("name")] return DbtRegularModel( @@ -145,6 +145,6 @@ def _regular_model_from_node(node: Dict[str, Any]) -> DbtRegularModel: ) -def regular_models_from_manifest(manifest: dict) -> List[DbtRegularModel]: +def regular_models_from_manifest(manifest: dict) -> list[DbtRegularModel]: """Turn the orphan nodes of a dbt manifest into ``DbtRegularModel`` instances.""" return [_regular_model_from_node(node) for node in find_orphan_model_nodes(manifest)] diff --git a/slayer/dbt/models.py b/slayer/dbt/models.py index d0e30b12..4623b8da 100644 --- a/slayer/dbt/models.py +++ b/slayer/dbt/models.py @@ -5,40 +5,39 @@ compatibility shim and has heavy transitive dependencies we don't need. """ -from typing import List, Optional from pydantic import BaseModel, Field, field_validator class DbtTimeTypeParams(BaseModel): - time_granularity: Optional[str] = None - is_partition: Optional[bool] = None + time_granularity: str | None = None + is_partition: bool | None = None class DbtNonAdditiveDimension(BaseModel): name: str window_choice: str = "min" - window_groupings: List[str] = Field(default_factory=list) + window_groupings: list[str] = Field(default_factory=list) class DbtEntity(BaseModel): name: str type: str # "primary", "foreign", "unique", "natural" - expr: Optional[str] = None # defaults to name if omitted - description: Optional[str] = None + expr: str | None = None # defaults to name if omitted + description: str | None = None class DbtDimension(BaseModel): name: str type: str = "categorical" # "categorical" or "time" - expr: Optional[str] = None - description: Optional[str] = None - label: Optional[str] = None - type_params: Optional[DbtTimeTypeParams] = None + expr: str | None = None + description: str | None = None + label: str | None = None + type_params: DbtTimeTypeParams | None = None class DbtMeasureAggParams(BaseModel): - percentile: Optional[float] = None + percentile: float | None = None use_discrete_percentile: bool = False use_approximate_percentile: bool = False @@ -46,17 +45,17 @@ class DbtMeasureAggParams(BaseModel): class DbtMeasure(BaseModel): name: str agg: str # "sum", "count", "average", "count_distinct", "min", "max", etc. - expr: Optional[str] = None - description: Optional[str] = None - label: Optional[str] = None - create_metric: Optional[bool] = None - agg_time_dimension: Optional[str] = None - agg_params: Optional[DbtMeasureAggParams] = None - non_additive_dimension: Optional[DbtNonAdditiveDimension] = None + expr: str | None = None + description: str | None = None + label: str | None = None + create_metric: bool | None = None + agg_time_dimension: str | None = None + agg_params: DbtMeasureAggParams | None = None + non_additive_dimension: DbtNonAdditiveDimension | None = None @field_validator("expr", mode="before") @classmethod - def _coerce_expr_to_str(cls, v: object) -> Optional[str]: + def _coerce_expr_to_str(cls, v: object) -> str | None: """Coerce numeric expr values to strings (e.g. dbt `expr: 1`).""" if v is None: return None @@ -64,60 +63,60 @@ def _coerce_expr_to_str(cls, v: object) -> Optional[str]: class DbtDefaults(BaseModel): - agg_time_dimension: Optional[str] = None + agg_time_dimension: str | None = None class DbtSemanticModel(BaseModel): name: str - model: Optional[str] = None # raw string, e.g. "ref('claim')" - description: Optional[str] = None - defaults: Optional[DbtDefaults] = None - primary_entity: Optional[str] = None - entities: List[DbtEntity] = Field(default_factory=list) - dimensions: List[DbtDimension] = Field(default_factory=list) - measures: List[DbtMeasure] = Field(default_factory=list) - label: Optional[str] = None + model: str | None = None # raw string, e.g. "ref('claim')" + description: str | None = None + defaults: DbtDefaults | None = None + primary_entity: str | None = None + entities: list[DbtEntity] = Field(default_factory=list) + dimensions: list[DbtDimension] = Field(default_factory=list) + measures: list[DbtMeasure] = Field(default_factory=list) + label: str | None = None class DbtMetricInputMeasure(BaseModel): """A measure reference within a metric's type_params.""" name: str - filter: Optional[str] = None - alias: Optional[str] = None + filter: str | None = None + alias: str | None = None class DbtMetricInput(BaseModel): """A metric reference within a derived metric's type_params.""" name: str - alias: Optional[str] = None - offset_window: Optional[str] = None - offset_to_grain: Optional[str] = None - filter: Optional[str] = None + alias: str | None = None + offset_window: str | None = None + offset_to_grain: str | None = None + filter: str | None = None class DbtMetricTypeParams(BaseModel): - measure: Optional[str] = None # simple metrics: measure name (string shorthand) - expr: Optional[str] = None # derived metrics: formula expression - metrics: Optional[List[DbtMetricInput]] = None # derived: input metric refs - numerator: Optional[DbtMetricInput] = None # ratio - denominator: Optional[DbtMetricInput] = None # ratio + measure: str | None = None # simple metrics: measure name (string shorthand) + expr: str | None = None # derived metrics: formula expression + metrics: list[DbtMetricInput] | None = None # derived: input metric refs + numerator: DbtMetricInput | None = None # ratio + denominator: DbtMetricInput | None = None # ratio class DbtMetric(BaseModel): name: str type: str # "simple", "derived", "cumulative", "ratio", "conversion" - description: Optional[str] = None - label: Optional[str] = None - type_params: Optional[DbtMetricTypeParams] = None - filter: Optional[str] = None + description: str | None = None + label: str | None = None + type_params: DbtMetricTypeParams | None = None + filter: str | None = None class DbtColumnMeta(BaseModel): """Column-level metadata from dbt's manifest for a regular (non-semantic) model.""" name: str - description: Optional[str] = None - data_type: Optional[str] = None - tags: List[str] = Field(default_factory=list) + description: str | None = None + data_type: str | None = None + tags: list[str] = Field(default_factory=list) class DbtRegularModel(BaseModel): @@ -137,17 +136,17 @@ class DbtRegularModel(BaseModel): regular model's SQL into a semantic-model-derived ``SlayerModel``. """ name: str - database: Optional[str] = None - schema_name: Optional[str] = None # avoids shadowing pydantic's `schema` method - alias: Optional[str] = None # materialized table name; falls back to `name` - description: Optional[str] = None - tags: List[str] = Field(default_factory=list) - columns: List[DbtColumnMeta] = Field(default_factory=list) - raw_code: Optional[str] = None # SQL body from the .sql file on disk, Jinja unresolved + database: str | None = None + schema_name: str | None = None # avoids shadowing pydantic's `schema` method + alias: str | None = None # materialized table name; falls back to `name` + description: str | None = None + tags: list[str] = Field(default_factory=list) + columns: list[DbtColumnMeta] = Field(default_factory=list) + raw_code: str | None = None # SQL body from the .sql file on disk, Jinja unresolved class DbtProject(BaseModel): """Aggregated result of parsing all YAML files in a dbt project.""" - semantic_models: List[DbtSemanticModel] = Field(default_factory=list) - metrics: List[DbtMetric] = Field(default_factory=list) - regular_models: List[DbtRegularModel] = Field(default_factory=list) + semantic_models: list[DbtSemanticModel] = Field(default_factory=list) + metrics: list[DbtMetric] = Field(default_factory=list) + regular_models: list[DbtRegularModel] = Field(default_factory=list) diff --git a/slayer/dbt/parser.py b/slayer/dbt/parser.py index fe324ad0..4a5c3903 100644 --- a/slayer/dbt/parser.py +++ b/slayer/dbt/parser.py @@ -13,7 +13,6 @@ import os import re from pathlib import Path -from typing import Dict, List import yaml @@ -56,7 +55,7 @@ def _extract_ref_name(raw: str) -> str: return raw -def _collect_yaml_paths(directory: str) -> List[str]: +def _collect_yaml_paths(directory: str) -> list[str]: """Recursively collect .yaml and .yml file paths, skipping hidden dirs/files.""" paths = [] for root, dirs, files in os.walk(directory): @@ -69,7 +68,7 @@ def _collect_yaml_paths(directory: str) -> List[str]: return paths -def _collect_sql_files(directory: str) -> Dict[str, str]: +def _collect_sql_files(directory: str) -> dict[str, str]: """Recursively collect .sql file bodies keyed by filename stem. dbt models are named after their `.sql` filename (without the extension), @@ -77,7 +76,7 @@ def _collect_sql_files(directory: str) -> Dict[str, str]: hidden dirs/files and any target/build directories dbt may have left behind. """ - result: Dict[str, str] = {} + result: dict[str, str] = {} for root, dirs, files in os.walk(directory): dirs[:] = [d for d in dirs if not d.startswith(".") and d != "target"] for filename in sorted(files): @@ -121,8 +120,8 @@ def parse_dbt_project( models whose underlying dbt model is a query rather than a physical table. """ - all_semantic_models: List[DbtSemanticModel] = [] - all_metrics: List[DbtMetric] = [] + all_semantic_models: list[DbtSemanticModel] = [] + all_metrics: list[DbtMetric] = [] yaml_paths = _collect_yaml_paths(project_path) if not yaml_paths: @@ -206,7 +205,7 @@ def parse_dbt_project( ) -def _parse_regular_models(project_path: str) -> List[DbtRegularModel]: +def _parse_regular_models(project_path: str) -> list[DbtRegularModel]: """Discover regular (non-semantic) dbt models via the dbt manifest. Returns an empty list when the manifest is absent and dbt-core is not diff --git a/slayer/dbt/sql_resolver.py b/slayer/dbt/sql_resolver.py index 3285833e..9603b5b7 100644 --- a/slayer/dbt/sql_resolver.py +++ b/slayer/dbt/sql_resolver.py @@ -24,7 +24,6 @@ import logging import re -from typing import Dict, List, Optional, Set, Tuple logger = logging.getLogger(__name__) @@ -68,11 +67,11 @@ def resolve_refs( sql: str, - regular_models_sql: Dict[str, str], + regular_models_sql: dict[str, str], *, max_depth: int = 16, - _visited: Optional[Set[str]] = None, -) -> Tuple[str, List[str]]: + _visited: set[str] | None = None, +) -> tuple[str, list[str]]: """Resolve dbt Jinja refs/sources in a SQL body. Refs pointing at other regular models are recursively inlined as @@ -99,7 +98,7 @@ def resolve_refs( if _visited is None: _visited = set() - warnings: List[str] = [] + warnings: list[str] = [] if max_depth <= 0: warnings.append( diff --git a/slayer/demo/jaffle_shop.py b/slayer/demo/jaffle_shop.py index 350b6bc9..255892d1 100644 --- a/slayer/demo/jaffle_shop.py +++ b/slayer/demo/jaffle_shop.py @@ -20,7 +20,7 @@ import tempfile from collections import deque from importlib.util import find_spec -from typing import IO, TYPE_CHECKING, List, Optional, Tuple +from typing import IO, TYPE_CHECKING from slayer.async_utils import run_sync from slayer.core.models import DatasourceConfig, SlayerModel @@ -129,7 +129,7 @@ def resolve_demo_db_path(storage_path: str) -> str: return os.path.join(demo_dir, "jaffle_shop.duckdb") -def _stream_fileno(stream) -> Optional[int]: +def _stream_fileno(stream) -> int | None: """Return ``stream.fileno()`` if it points at a real file descriptor, else None. ipykernel / nbclient replace ``sys.stdout`` / ``sys.stderr`` with shim @@ -144,7 +144,7 @@ def _stream_fileno(stream) -> Optional[int]: return None -def _jafgen_cmd(years: int) -> List[str]: +def _jafgen_cmd(years: int) -> list[str]: """Build the jafgen invocation without relying on PATH exposure. ``uv tool install motley-slayer`` (and pipx) link only slayer's own entry @@ -171,7 +171,7 @@ def generate_data( output_dir: str, years: int = 1, *, - stream: Optional[IO[str]] = None, + stream: IO[str] | None = None, ) -> str: """Run ``jafgen`` into ``output_dir``; return the path to the generated CSVs. @@ -232,7 +232,7 @@ def generate_data( def create_schema( conn: "duckdb.DuckDBPyConnection", - schema_path: Optional[str] = None, + schema_path: str | None = None, ) -> None: """Create the Jaffle Shop tables on ``conn``. @@ -372,7 +372,7 @@ def build_jaffle_shop( *, years: int = 2, force: bool = False, - stream: Optional[IO[str]] = None, + stream: IO[str] | None = None, ) -> bool: """Generate the Jaffle Shop DuckDB at ``db_path`` if it does not already exist. @@ -415,8 +415,8 @@ def ensure_demo_datasource( years: int = 2, ingest_models: bool = True, assume_yes: bool = True, - stream: Optional[IO[str]] = None, -) -> Tuple[DatasourceConfig, List[SlayerModel], bool]: + stream: IO[str] | None = None, +) -> tuple[DatasourceConfig, list[SlayerModel], bool]: """Ensure the Jaffle Shop demo is fully set up in ``storage``. - Builds the DuckDB at ``/demo/jaffle_shop.duckdb`` if missing. @@ -465,11 +465,11 @@ def ensure_demo_datasource( from slayer.engine.ingestion import ingest_datasource models = ingest_datasource(datasource=ds) - written: List[SlayerModel] = [] + written: list[SlayerModel] = [] for model in models: if model.name in DEFAULT_TIME_DIMENSIONS: model.default_time_dimension = DEFAULT_TIME_DIMENSIONS[model.name] - existing_model: Optional[SlayerModel] = run_sync( + existing_model: SlayerModel | None = run_sync( storage.get_model(name=model.name, data_source=name) ) if existing_model is not None and not assume_yes: diff --git a/slayer/embeddings/client.py b/slayer/embeddings/client.py index 0404da74..518e108e 100644 --- a/slayer/embeddings/client.py +++ b/slayer/embeddings/client.py @@ -25,7 +25,7 @@ import os import warnings from functools import lru_cache -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import Any DEFAULT_EMBEDDING_MODEL = "openai/text-embedding-3-small" @@ -109,10 +109,10 @@ def is_available() -> bool: _HASH_PREFIX_CHARS = 16 -_CAP_CACHE: Dict[str, int] = {} +_CAP_CACHE: dict[str, int] = {} -def _try_get_max_tokens(model: str) -> Optional[int]: +def _try_get_max_tokens(model: str) -> int | None: """Best-effort ``litellm.utils.get_max_tokens`` lookup. Returns the cap if it's a positive int, else ``None``. Never raises.""" try: @@ -128,7 +128,7 @@ def _try_get_max_tokens(model: str) -> Optional[int]: return None -def _resolve_model_cap(resolved_model: str) -> Optional[int]: +def _resolve_model_cap(resolved_model: str) -> int | None: """Look up the model's max-token cap, trying provider-prefixed name first and the bare name as a fallback. Caches only positive-int successes so a transient lookup failure isn't sticky.""" @@ -173,7 +173,7 @@ def _strip_provider_prefix(model: str) -> str: def truncate_text_for_model( - text: str, *, model: Optional[str] = None, + text: str, *, model: str | None = None, ) -> str: """Truncate ``text`` to the resolved model's token cap minus a fixed 256-token margin (DEV-1557). @@ -230,14 +230,14 @@ def truncate_text_for_model( return truncated -def _get_bad_request_exception_classes() -> Tuple[Type[BaseException], ...]: +def _get_bad_request_exception_classes() -> tuple[type[BaseException], ...]: """Return the exception class(es) representing litellm's ``BadRequestError`` — empty tuple if neither ``litellm`` nor ``litellm.exceptions`` exposes one. Tuple form lets us use the result directly in an ``except`` clause; an empty tuple safely catches nothing so the generic-exception fallback path takes over. """ - classes: List[Type[BaseException]] = [] + classes: list[type[BaseException]] = [] try: import litellm # noqa: PLC0415 — lazy except Exception: # noqa: BLE001 @@ -260,11 +260,11 @@ def _get_bad_request_exception_classes() -> Tuple[Type[BaseException], ...]: return tuple(classes) -def _parse_vectors(response: Any, n: int) -> List[Optional[List[float]]]: +def _parse_vectors(response: Any, n: int) -> list[list[float] | None]: """Pack a litellm aembedding response into ``n`` per-slot vectors, padding short responses with ``None``.""" data = getattr(response, "data", None) or [] - out: List[Optional[List[float]]] = [] + out: list[list[float] | None] = [] for entry in data: if isinstance(entry, dict): vec = entry.get("embedding") @@ -282,15 +282,15 @@ def _parse_vectors(response: Any, n: int) -> List[Optional[List[float]]]: async def _per_input_retry( litellm: Any, resolved_model: str, - truncated_texts: List[str], - bad_request_classes: Tuple[Type[BaseException], ...], -) -> List[Optional[List[float]]]: + truncated_texts: list[str], + bad_request_classes: tuple[type[BaseException], ...], +) -> list[list[float] | None]: """Embed each text in ``truncated_texts`` individually. On a per-text ``BadRequestError`` the slot is ``None`` and the loop continues; on any other exception we treat it as a global failure shape (rate limit / auth / network), mark the current slot and every remaining slot ``None``, and return early.""" - results: List[Optional[List[float]]] = [] + results: list[list[float] | None] = [] for idx, text in enumerate(truncated_texts): try: response = await litellm.aembedding( @@ -320,8 +320,8 @@ async def _per_input_retry( async def embed_batch( - texts: List[str], *, model: Optional[str] = None, -) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, +) -> list[list[float] | None]: """Embed a batch of texts via ``litellm.aembedding`` (DEV-1557). Each text is preemptively truncated to fit the resolved model's @@ -372,11 +372,11 @@ async def embed_batch( return _parse_vectors(response, len(truncated)) -_QUERY_CACHE: "dict[tuple[str, str], List[float]]" = {} +_QUERY_CACHE: "dict[tuple[str, str], list[float]]" = {} _QUERY_CACHE_MAX = 64 -async def embed_query(text: str, *, model: Optional[str] = None) -> Optional[List[float]]: +async def embed_query(text: str, *, model: str | None = None) -> list[float] | None: """Embed a single query string with a small process-wide LRU cache. Returns ``None`` when the extra is not installed or the embedding call diff --git a/slayer/embeddings/models.py b/slayer/embeddings/models.py index 4110fc15..bf534914 100644 --- a/slayer/embeddings/models.py +++ b/slayer/embeddings/models.py @@ -18,7 +18,7 @@ """ from datetime import datetime, timezone -from typing import Any, List, Literal +from typing import Any, Literal from pydantic import BaseModel, Field, model_validator @@ -42,7 +42,7 @@ class Embedding(BaseModel): embedding_model_name: str entity_kind: EntityKind content_hash: str - embedding: List[float] + embedding: list[float] created_at: datetime = Field(default_factory=_utcnow) @model_validator(mode="before") diff --git a/slayer/embeddings/ranker.py b/slayer/embeddings/ranker.py index dca9af89..49055933 100644 --- a/slayer/embeddings/ranker.py +++ b/slayer/embeddings/ranker.py @@ -13,12 +13,11 @@ from __future__ import annotations -from typing import List, Tuple import numpy as np -def normalise(vector: List[float]) -> np.ndarray: +def normalise(vector: list[float]) -> np.ndarray: """Return a unit-L2 numpy view of ``vector``. Zero vectors come back unchanged so we never divide by zero.""" arr = np.asarray(vector, dtype=np.float32) @@ -42,7 +41,7 @@ def top_k_cosine( query: np.ndarray, matrix: np.ndarray, k: int, -) -> List[Tuple[int, float]]: +) -> list[tuple[int, float]]: """Return the top-``k`` ``(row_index, cosine_similarity)`` pairs. Assumes ``query`` is already unit-normalised (1D shape ``(dim,)``) and diff --git a/slayer/engine/column_dependency.py b/slayer/engine/column_dependency.py index c3cfaf99..29d282dd 100644 --- a/slayer/engine/column_dependency.py +++ b/slayer/engine/column_dependency.py @@ -18,7 +18,7 @@ from __future__ import annotations from collections import deque -from typing import TYPE_CHECKING, Deque, Dict, List, Optional, Set, Tuple +from typing import TYPE_CHECKING import sqlglot from sqlglot import exp @@ -35,15 +35,15 @@ # ``exp.Column`` identifier shape — dialect choice does not change which # columns appear in the AST. Using sqlglot's default keeps the validator # independent of the model's runtime datasource dialect. -_DEPENDENCY_DIALECT: Optional[str] = None +_DEPENDENCY_DIALECT: str | None = None def _resolve_target_for_ref( *, - table_alias: Optional[str], + table_alias: str | None, host: SlayerModel, - reachable: Dict[str, SlayerModel], -) -> Optional[SlayerModel]: + reachable: dict[str, SlayerModel], +) -> SlayerModel | None: """Return the model that a column reference resolves to, or ``None``. Mirrors the runtime alias resolution in @@ -85,9 +85,9 @@ def _resolve_single_column( *, node: exp.Column, host: SlayerModel, - reachable: Dict[str, SlayerModel], - root_ids: Set[int], -) -> Optional[Tuple[str, str]]: + reachable: dict[str, SlayerModel], + root_ids: set[int], +) -> tuple[str, str] | None: """Resolve one ``exp.Column`` node to a (model_name, column_name) dep, or ``None`` if the node is out of scope. @@ -116,8 +116,8 @@ def _column_dependencies( *, column: Column, host: SlayerModel, - reachable: Dict[str, SlayerModel], -) -> List[Tuple[str, str]]: + reachable: dict[str, SlayerModel], +) -> list[tuple[str, str]]: """Extract the root-scope derived-column dependencies of ``column``. Returns a list of ``(model_name, column_name)`` tuples — only refs @@ -135,7 +135,7 @@ def _column_dependencies( # sees, not a noisy validator complaint about unparseable SQL. return [] root_ids = _root_scope_column_ids(parsed=parsed) - deps: List[Tuple[str, str]] = [] + deps: list[tuple[str, str]] = [] for node in parsed.find_all(exp.Column): resolved = _resolve_single_column( node=node, host=host, reachable=reachable, root_ids=root_ids, @@ -147,9 +147,9 @@ def _column_dependencies( def _node_dependencies( *, - node: Tuple[str, str], - reachable: Dict[str, SlayerModel], -) -> List[Tuple[str, str]]: + node: tuple[str, str], + reachable: dict[str, SlayerModel], +) -> list[tuple[str, str]]: """Return the dependency edges leaving ``node = (model_name, col_name)``. Empty list when the model or column is missing — those are dead-ends, not errors. @@ -166,12 +166,12 @@ def _node_dependencies( def _dfs_visit( *, - node: Tuple[str, str], - reachable: Dict[str, SlayerModel], - on_stack: List[Tuple[str, str]], - on_stack_set: Set[Tuple[str, str]], - visited: Set[Tuple[str, str]], -) -> Optional[List[Tuple[str, str]]]: + node: tuple[str, str], + reachable: dict[str, SlayerModel], + on_stack: list[tuple[str, str]], + on_stack_set: set[tuple[str, str]], + visited: set[tuple[str, str]], +) -> list[tuple[str, str]] | None: """Recursive DFS visit. Returns the first cycle reachable from ``node``, or ``None``. Mutates ``on_stack`` / ``on_stack_set`` / ``visited`` in place — the caller initialises them empty and @@ -199,9 +199,9 @@ def _dfs_visit( def _detect_cycle_dfs( *, - start: Tuple[str, str], - reachable: Dict[str, SlayerModel], -) -> Optional[List[Tuple[str, str]]]: + start: tuple[str, str], + reachable: dict[str, SlayerModel], +) -> list[tuple[str, str]] | None: """DFS from ``start = (model_name, column_name)``. Returns the first cycle found as an ordered list (start may appear at both ends if the cycle closes through it), or ``None`` if the subgraph is acyclic. @@ -216,14 +216,14 @@ async def _prefetch_reachable_models( *, model: SlayerModel, storage: "StorageBackend", -) -> Dict[str, SlayerModel]: +) -> dict[str, SlayerModel]: """BFS over ``model.joins`` (transitively), pulling each target model in the same ``data_source``. Returns ``{model_name: model}`` including ``model`` itself. Unresolvable target names (model not persisted yet) are silently omitted — save-time is best-effort. """ - out: Dict[str, SlayerModel] = {model.name: model} - queue: Deque[SlayerModel] = deque([model]) + out: dict[str, SlayerModel] = {model.name: model} + queue: deque[SlayerModel] = deque([model]) while queue: current = queue.popleft() for join in current.joins: @@ -259,7 +259,7 @@ async def validate_no_column_cycles( reachable = await _prefetch_reachable_models(model=model, storage=storage) # Iterate roots in a deterministic order so the reported cycle is # stable across runs. - roots: List[Tuple[str, str]] = [] + roots: list[tuple[str, str]] = [] for entity_name in sorted(reachable.keys()): entity = reachable[entity_name] for col in entity.columns: diff --git a/slayer/engine/column_expansion.py b/slayer/engine/column_expansion.py index aa20546b..2cd45285 100644 --- a/slayer/engine/column_expansion.py +++ b/slayer/engine/column_expansion.py @@ -17,7 +17,8 @@ """ from __future__ import annotations -from typing import Any, Awaitable, Callable, Dict, Optional, Set, Tuple +from typing import Any +from collections.abc import Awaitable, Callable import sqlglot from sqlglot import exp @@ -26,7 +27,7 @@ from slayer.core.errors import ColumnCycleError from slayer.core.models import Column, SlayerModel -ResolveModel = Callable[..., Awaitable[Optional[SlayerModel]]] +ResolveModel = Callable[..., Awaitable[SlayerModel | None]] def _is_trivial_base(*, column: Column) -> bool: @@ -46,7 +47,7 @@ def _is_trivial_base(*, column: Column) -> bool: return sql == column.name -def _root_scope_column_ids(*, parsed: exp.Expression) -> Set[int]: +def _root_scope_column_ids(*, parsed: exp.Expression) -> set[int]: """Return the ``id()`` set of ``exp.Column`` nodes that lexically belong to the root scope of ``parsed`` (DEV-1410). @@ -67,7 +68,7 @@ def _root_scope_column_ids(*, parsed: exp.Expression) -> Set[int]: if not isinstance(parsed, exp.Expression): return set() wrapper = exp.Select(expressions=[exp.Alias(this=parsed.copy(), alias="_")]) - scope_node_ids: Dict[int, ScopeType] = {} + scope_node_ids: dict[int, ScopeType] = {} for scope in traverse_scope(wrapper): scope_node_ids[id(scope.expression)] = scope.scope_type if not scope_node_ids: @@ -89,10 +90,10 @@ def _root_scope_column_ids(*, parsed: exp.Expression) -> Set[int]: # wrapper just wraps a deep copy and ``find_all`` walks in # document order. return set() - root_ids: Set[int] = set() + root_ids: set[int] = set() for w_col, p_col in zip(wrapper_cols, parsed_cols): - node: Optional[exp.Expression] = w_col.parent - scope_type: Optional[ScopeType] = None + node: exp.Expression | None = w_col.parent + scope_type: ScopeType | None = None while node is not None: if id(node) in scope_node_ids: scope_type = scope_node_ids[id(node)] @@ -109,9 +110,9 @@ async def _walk_path_to_target( source_alias: str, table_alias: str, resolve_model: ResolveModel, - named_queries: Dict[str, Any], + named_queries: dict[str, Any], is_root: bool, -) -> Tuple[Optional[SlayerModel], Optional[str]]: +) -> tuple[SlayerModel | None, str | None]: """Resolve a ``table_alias`` (e.g. ``B`` or ``B__C``) seen inside a Column.sql to the terminal joined model and the canonical alias to use in emitted SQL. @@ -157,11 +158,11 @@ async def _process_column_node( model: SlayerModel, alias_path: str, resolve_model: ResolveModel, - named_queries: Dict[str, Any], + named_queries: dict[str, Any], dialect: str, - visited: Tuple[Tuple[str, str], ...], + visited: tuple[tuple[str, str], ...], is_root: bool, - root_scope_ids: Set[int], + root_scope_ids: set[int], ) -> None: """Resolve one ``exp.Column`` node in the parsed AST, mutating it in place. Encapsulates the multi-branch decision that drives expansion: @@ -248,15 +249,15 @@ async def _process_column_node( async def expand_derived_refs( *, - sql: Optional[str], + sql: str | None, model: SlayerModel, alias_path: str, resolve_model: ResolveModel, - named_queries: Optional[Dict[str, Any]] = None, + named_queries: dict[str, Any] | None = None, dialect: str, - visited: Optional[Tuple[Tuple[str, str], ...]] = None, + visited: tuple[tuple[str, str], ...] | None = None, is_root: bool = True, -) -> Optional[str]: +) -> str | None: """Recursively expand cross-model and local derived-column references inside ``sql``. diff --git a/slayer/engine/enriched.py b/slayer/engine/enriched.py index 54ddb2c1..a6ef17c1 100644 --- a/slayer/engine/enriched.py +++ b/slayer/engine/enriched.py @@ -17,7 +17,7 @@ - The query engine controls resolution logic (placeholder expansion, join resolution) """ -from typing import Dict, List, Optional +from typing import Optional from pydantic import BaseModel, Field @@ -32,19 +32,19 @@ class EnrichedDimension(BaseModel): """A dimension with its SQL expression fully resolved.""" name: str - sql: Optional[str] + sql: str | None type: DataType alias: str = Field(description="Result column name, e.g. 'orders.status'") model_name: str - label: Optional[str] = Field(default=None, description="Human-readable label") - format: Optional[NumberFormat] = Field(default=None, description="Number format from the source dimension") + label: str | None = Field(default=None, description="Human-readable label") + format: NumberFormat | None = Field(default=None, description="Number format from the source dimension") class EnrichedMeasure(BaseModel): """A measure with its SQL expression and aggregation fully resolved.""" name: str - sql: Optional[str] = Field(description="SQL expression; None for *:count (COUNT(*))") + sql: str | None = Field(description="SQL expression; None for *:count (COUNT(*))") aggregation: str = Field(description="Aggregation name: sum, avg, count, weighted_avg, etc.") alias: str = Field(description="Result column name, e.g. 'orders.revenue_sum'") user_declared: bool = Field( @@ -57,19 +57,19 @@ class EnrichedMeasure(BaseModel): ), ) model_name: str - aggregation_def: Optional[Aggregation] = Field(default=None, description="Full aggregation definition (formula, params)") - agg_kwargs: Dict[str, str] = Field(default_factory=dict, description="Query-time aggregation param overrides") - window: Optional[str] = Field(default=None, description="Trailing time window for windowed sum/avg aggregations") - window_time_alias: Optional[str] = Field(default=None, description="Time dimension alias used for windowed aggregations") - label: Optional[str] = Field(default=None, description="Human-readable label") - time_column: Optional[str] = Field(default=None, description="Explicit time col for first/last (overrides query default)") - source_measure_name: Optional[str] = Field(default=None, description="Original measure name before canonicalization") - filter_sql: Optional[str] = Field(default=None, description="Resolved SQL condition for filtered measures (CASE WHEN)") - filter_columns: List[str] = Field( + aggregation_def: Aggregation | None = Field(default=None, description="Full aggregation definition (formula, params)") + agg_kwargs: dict[str, str] = Field(default_factory=dict, description="Query-time aggregation param overrides") + window: str | None = Field(default=None, description="Trailing time window for windowed sum/avg aggregations") + window_time_alias: str | None = Field(default=None, description="Time dimension alias used for windowed aggregations") + label: str | None = Field(default=None, description="Human-readable label") + time_column: str | None = Field(default=None, description="Explicit time col for first/last (overrides query default)") + source_measure_name: str | None = Field(default=None, description="Original measure name before canonicalization") + filter_sql: str | None = Field(default=None, description="Resolved SQL condition for filtered measures (CASE WHEN)") + filter_columns: list[str] = Field( default_factory=list, description="Resolved (qualified) column names referenced by the filter, for join planning", ) - type: Optional[DataType] = Field( + type: DataType | None = Field( default=None, description=( "DEV-1361: declared result type of the aggregation. When set, " @@ -77,7 +77,7 @@ class EnrichedMeasure(BaseModel): "Inherits from ModelMeasure.type at enrichment time." ), ) - column_type: Optional[DataType] = Field( + column_type: DataType | None = Field( default=None, description=( "DEV-1361: source column's declared type — wraps the inner " @@ -104,12 +104,12 @@ class EnrichedTimeDimension(BaseModel): """A time dimension with resolved SQL and granularity.""" name: str - sql: Optional[str] + sql: str | None granularity: TimeGranularity - date_range: Optional[List[str]] + date_range: list[str] | None alias: str model_name: str - label: Optional[str] = None + label: str | None = None class EnrichedExpression(BaseModel): @@ -130,8 +130,8 @@ class EnrichedExpression(BaseModel): "expressions (e.g. desugared change/change_pct arithmetic)." ), ) - label: Optional[str] = None - type: Optional[DataType] = Field( + label: str | None = None + type: DataType | None = Field( default=None, description=( "DEV-1361: declared result type — when set, the outer SELECT " @@ -163,10 +163,10 @@ class EnrichedTransform(BaseModel): measure_alias: str = Field(description="Alias of the measure in the base CTE to transform") alias: str = Field(description="Result column name") offset: int = Field(description="For time_shift: number of rows or calendar units") - granularity: Optional[str] = Field(default=None, description="For time_shift: year, month, quarter, etc.") - time_alias: Optional[str] = Field(default=None, description="Alias of the time dimension column for ORDER BY") - partition_aliases: List[str] = Field(default_factory=list, description="Dimension aliases to PARTITION BY") - n: Optional[int] = Field(default=None, description="Bucket count for ntile(measure, n=...)") + granularity: str | None = Field(default=None, description="For time_shift: year, month, quarter, etc.") + time_alias: str | None = Field(default=None, description="Alias of the time dimension column for ORDER BY") + partition_aliases: list[str] = Field(default_factory=list, description="Dimension aliases to PARTITION BY") + n: int | None = Field(default=None, description="Bucket count for ntile(measure, n=...)") predicate_is_boolean: bool = Field( default=False, description="True when the transform's measure_alias points at a boolean expression " @@ -174,8 +174,8 @@ class EnrichedTransform(BaseModel): "Postgres rejects 'boolean <> integer' so the numeric `IS NOT NULL AND <> 0` " "predicate cannot be used.", ) - label: Optional[str] = None - type: Optional[DataType] = Field( + label: str | None = None + type: DataType | None = Field( default=None, description=( "DEV-1361: declared result type — when set, the window-layer " @@ -192,30 +192,30 @@ class EnrichedQuery(BaseModel): """ model_name: str - sql_table: Optional[str] = None - sql: Optional[str] = None + sql_table: str | None = None + sql: str | None = None - resolved_joins: List[tuple] = Field(default_factory=list, description="[(target_table_sql, target_alias, join_condition, join_type), ...]") + resolved_joins: list[tuple] = Field(default_factory=list, description="[(target_table_sql, target_alias, join_condition, join_type), ...]") - dimensions: List[EnrichedDimension] = Field(default_factory=list) - measures: List[EnrichedMeasure] = Field(default_factory=list) - time_dimensions: List[EnrichedTimeDimension] = Field(default_factory=list) + dimensions: list[EnrichedDimension] = Field(default_factory=list) + measures: list[EnrichedMeasure] = Field(default_factory=list) + time_dimensions: list[EnrichedTimeDimension] = Field(default_factory=list) - expressions: List[EnrichedExpression] = Field(default_factory=list) - transforms: List[EnrichedTransform] = Field(default_factory=list) + expressions: list[EnrichedExpression] = Field(default_factory=list) + transforms: list[EnrichedTransform] = Field(default_factory=list) - cross_model_measures: List["CrossModelMeasure"] = Field(default_factory=list) + cross_model_measures: list["CrossModelMeasure"] = Field(default_factory=list) - last_agg_time_column: Optional[str] = Field(default=None, description="Time column for first/last aggregation (ORDER BY for ROW_NUMBER)") + last_agg_time_column: str | None = Field(default=None, description="Time column for first/last aggregation (ORDER BY for ROW_NUMBER)") - filters: List[ParsedFilter] = Field(default_factory=list) - order: Optional[List[OrderItem]] = None - limit: Optional[int] = None - offset: Optional[int] = None + filters: list[ParsedFilter] = Field(default_factory=list) + order: list[OrderItem] | None = None + limit: int | None = None + offset: int | None = None - field_name_aliases: Dict[str, str] = Field(default_factory=dict, description="Custom field name → enriched alias mapping (for ORDER BY resolution)") + field_name_aliases: dict[str, str] = Field(default_factory=dict, description="Custom field name → enriched alias mapping (for ORDER BY resolution)") - user_projection: List[str] = Field( + user_projection: list[str] = Field( default_factory=list, description=( "DEV-1444: ordered list of result-column aliases that the user " @@ -254,18 +254,18 @@ class CrossModelMeasure(BaseModel): ) alias: str = Field(description="Result column name, e.g. 'orders.customers__avg_score'") target_model_name: str = Field(description="The joined model name") - target_model_sql_table: Optional[str] - target_model_sql: Optional[str] + target_model_sql_table: str | None + target_model_sql: str | None measure: EnrichedMeasure = Field(description="The measure to aggregate") - join_pairs: List[List[str]] = Field(description="[[source_dim, target_dim], ...] from ModelJoin") - shared_dimensions: List[EnrichedDimension] = Field(description="Dimensions shared between main and target") - shared_time_dimensions: List[EnrichedTimeDimension] = Field(description="Time dims shared between main and target") + join_pairs: list[list[str]] = Field(description="[[source_dim, target_dim], ...] from ModelJoin") + shared_dimensions: list[EnrichedDimension] = Field(description="Dimensions shared between main and target") + shared_time_dimensions: list[EnrichedTimeDimension] = Field(description="Time dims shared between main and target") source_model_name: str = Field(description="The main query's model name") - source_sql_table: Optional[str] = Field(description="Main model's table") - source_sql: Optional[str] = Field(description="Main model's SQL") + source_sql_table: str | None = Field(description="Main model's table") + source_sql: str | None = Field(description="Main model's SQL") join_type: str = Field(default="left", description="'left' or 'inner'") - label: Optional[str] = None - format: Optional[NumberFormat] = Field(default=None, description="Inferred format for this cross-model measure") + label: str | None = None + format: NumberFormat | None = Field(default=None, description="Inferred format for this cross-model measure") rerooted_enriched: Optional["EnrichedQuery"] = Field(default=None, description="Re-rooted subquery with target as source") @@ -274,7 +274,7 @@ class CrossModelMeasure(BaseModel): CrossModelMeasure.model_rebuild() -def public_projection_aliases(enriched: EnrichedQuery) -> List[str]: +def public_projection_aliases(enriched: EnrichedQuery) -> list[str]: """Return the ordered list of public-projection aliases for ``enriched``. DEV-1444: the outer rendered SELECT projects exactly the @@ -301,7 +301,7 @@ def public_projection_aliases(enriched: EnrichedQuery) -> List[str]: # desugaring). Matches the pre-DEV-1444 ``expected_columns`` rule # in ``query_engine.py``. internal_prefixes = ("_inner_", "_ft", "_ts") - out: List[str] = [d.alias for d in enriched.dimensions] + out: list[str] = [d.alias for d in enriched.dimensions] out.extend(td.alias for td in enriched.time_dimensions) out.extend( m.alias for m in enriched.measures diff --git a/slayer/engine/enrichment.py b/slayer/engine/enrichment.py index effbf23b..48cb4d13 100644 --- a/slayer/engine/enrichment.py +++ b/slayer/engine/enrichment.py @@ -10,7 +10,8 @@ import difflib import re -from typing import Any, Dict, List, Mapping, Optional, Set, Tuple +from typing import Any +from collections.abc import Mapping import sqlglot from sqlglot import exp @@ -64,7 +65,7 @@ def _strip_string_literal(value: str) -> str: _canonical_agg_name = canonical_agg_name # Module-internal alias for the shared helper -def canonical_expression_key(node: Any) -> Tuple[Any, ...]: # NOSONAR(S8495) — variable-length tuple shape IS the discriminator; type signature already declares Tuple[Any, ...] +def canonical_expression_key(node: Any) -> tuple[Any, ...]: # NOSONAR(S8495) — variable-length tuple shape IS the discriminator; type signature already declares Tuple[Any, ...] """DEV-1444: build an alias-independent, structural hash key for a parsed formula AST node. @@ -111,8 +112,8 @@ def canonical_expression_key(node: Any) -> Tuple[Any, ...]: # NOSONAR(S8495) async def _collect_reachable_agg_names( model: SlayerModel, resolve_join_target, - named_queries: Dict, -) -> Optional[frozenset[str]]: + named_queries: dict, +) -> frozenset[str] | None: """Collect custom aggregation names from the source model and all reachable joined models. Walks the full reachable join graph via BFS, bounded only by the ``visited`` @@ -150,7 +151,7 @@ async def _collect_reachable_agg_names( async def enrich_query( query: SlayerQuery, model: SlayerModel, - named_queries: Optional[Dict[str, SlayerQuery]] = None, + named_queries: dict[str, SlayerQuery] | None = None, *, resolve_dimension_via_joins, resolve_cross_model_measure, @@ -194,7 +195,7 @@ async def enrich_query( # Saved-formula library for bare-name resolution. Only the source model's # named measures are in scope here; cross-model references (`other.aov`) # remain handled by the cross-model resolver. - named_measures: Dict[str, str] = {} + named_measures: dict[str, str] = {} for m in model.measures: if not m.name: continue @@ -231,7 +232,7 @@ async def enrich_query( ) # --- Measures (populated from fields below) --- - measures: List[EnrichedMeasure] = [] + measures: list[EnrichedMeasure] = [] # --- Time dimensions --- time_dimensions = await _resolve_time_dimensions( @@ -273,22 +274,22 @@ async def enrich_query( ) # --- Process fields --- - enriched_expressions: List[EnrichedExpression] = [] - enriched_transforms: List[EnrichedTransform] = [] - cross_model_measures: List[CrossModelMeasure] = [] - known_aliases: Dict[str, str] = {} - field_name_aliases: Dict[str, str] = {} + enriched_expressions: list[EnrichedExpression] = [] + enriched_transforms: list[EnrichedTransform] = [] + cross_model_measures: list[CrossModelMeasure] = [] + known_aliases: dict[str, str] = {} + field_name_aliases: dict[str, str] = {} # DEV-1443: canonical-agg alias → user-supplied measure name. Populated # when a query measure renames the canonical (``{"formula": "col:agg", # "name": "alias"}``). Consumed by the filter pre-pass (so a filter # written as ``col:agg N`` resolves to the user alias and HAVINGs # correctly) and by the ORDER BY enrichment (same shape). - canonical_to_user_name: Dict[str, str] = {} + canonical_to_user_name: dict[str, str] = {} # Cached source-column name set for the remap eligibility guard # (Codex Finding 1 — skip remap when the canonical alias also literally # names a source column on the model, since the regex sub would then # clobber the literal source-column reference). - _source_column_names: Set[str] = {c.name for c in model.columns} + _source_column_names: set[str] = {c.name for c in model.columns} # DEV-1444 provenance-merge index: canonical_expression_key → # surfaced alias. Populated when an EnrichedMeasure is created; @@ -297,7 +298,7 @@ async def enrich_query( # (e.g. order-by ``revenue:sum`` matching a user-declared # ``{"formula":"revenue:sum","name":"total"}``) reuses the existing # alias instead of materialising a phantom ``orders.revenue_sum``. - measure_canonical_key_to_alias: Dict[Tuple[Any, ...], str] = {} + measure_canonical_key_to_alias: dict[tuple[Any, ...], str] = {} def _mark_user_declared(alias: str) -> bool: """DEV-1444: flip ``user_declared=True`` on the enriched entry that @@ -327,8 +328,8 @@ async def _ensure_aggregated_measure( alias_key: str, measure_name: str, aggregation_name: str, - agg_args: Optional[list] = None, - agg_kwargs: Optional[dict] = None, + agg_args: list | None = None, + agg_kwargs: dict | None = None, ): """Create an EnrichedMeasure for an aggregated measure ref. @@ -498,7 +499,7 @@ async def _ensure_aggregated_measure( # syntax, transform calls, ``OVER``) were rejected at construction # by ``parse_sql_predicate``. filter_sql = None - filter_columns: List[str] = [] + filter_columns: list[str] = [] if measure_def and measure_def.filter: parsed = parse_sql_predicate(measure_def.filter) resolved = await resolve_filter_columns( @@ -555,7 +556,7 @@ def _resolve_sql(sql: str) -> str: resolved = re.sub(rf'(? List[str]: + def _resolve_rank_partition(transform: str, partition_by: list[str]) -> list[str]: """Resolve partition_by= column references to base-CTE aliases. partition_by entries must reference query dimensions or time dimensions — @@ -571,7 +572,7 @@ def _resolve_rank_partition(transform: str, partition_by: List[str]) -> List[str by_name.setdefault(td.name, td.alias) by_alias.setdefault(td.alias, td.alias) - resolved: List[str] = [] + resolved: list[str] = [] for col in partition_by: if col in by_alias: resolved.append(by_alias[col]) @@ -593,7 +594,7 @@ def _add_transform( offset: int = 1, granularity: str = None, predicate_is_boolean: bool = False, - kwargs: Optional[Dict[str, Any]] = None, + kwargs: dict[str, Any] | None = None, ): needs_time = transform in TIME_TRANSFORMS if needs_time and resolved_time_alias is None: @@ -644,7 +645,7 @@ def _add_transform( # falls through to the cross-model CTE path for them. _DISTRIBUTIVE_AGGS = frozenset({"sum", "min", "max"}) - def _intercept_candidate_for_cross_model(ref) -> "Optional[Tuple[str, str]]": + def _intercept_candidate_for_cross_model(ref) -> "tuple[str, str] | None": """DEV-1449: return ``(flat_with_agg, outer_agg)`` if the intercept would resolve a virtual-stage cross-model agg ref to a local re-aggregation on a flat column, or ``None`` if the @@ -707,7 +708,7 @@ def _intercept_candidate_for_cross_model(ref) -> "Optional[Tuple[str, str]]": async def _try_intercept_cross_model_as_local( ref, field_name: str, - ) -> Optional[str]: + ) -> str | None: """Apply the intercept (computes candidate + builds the EnrichedMeasure). Returns the full enriched alias the caller can use, or ``None`` if no candidate.""" @@ -749,7 +750,7 @@ async def _try_intercept_cross_model_as_local( field_name_aliases[ref_canonical] = local_alias return local_alias - async def _ensure_measure_from_spec(mname: str, agg_refs: Optional[dict] = None): + async def _ensure_measure_from_spec(mname: str, agg_refs: dict | None = None): """Ensure a measure is resolved — handles agg refs only.""" agg_refs = agg_refs or {} if mname in agg_refs: @@ -986,7 +987,7 @@ async def _flatten_spec(spec, field_name: str) -> str: # DEV-1444: track aliases in declared order so EnrichedQuery.user_projection # can be populated at the end. Dims and time-dims come first. - user_projection: List[str] = [d.alias for d in dimensions] + user_projection: list[str] = [d.alias for d in dimensions] user_projection.extend(td.alias for td in time_dimensions) # DEV-1444 (Codex review on PR #134): the provenance-merge index in @@ -999,7 +1000,7 @@ async def _flatten_spec(spec, field_name: str) -> str: # outer trim would then project a column the inner SELECT doesn't # expose. Track the set of canonical keys already owned by a # user-declared qfield and refuse the duplicate. - user_declared_canon_keys: Dict[Tuple[Any, ...], str] = {} + user_declared_canon_keys: dict[tuple[Any, ...], str] = {} # DEV-1443 (CodeRabbit thread + Codex round 4 on PR #133): the # duplicate-explicit-name check must run for every query measure @@ -1008,7 +1009,7 @@ async def _flatten_spec(spec, field_name: str) -> str: # arithmetic/transform measures fall through to ``_flatten_spec``; in # both cases a duplicate ``name`` would silently collapse two # measures onto a single alias. Run the pairwise check once up front. - _seen_explicit_names: Dict[str, str] = {} + _seen_explicit_names: dict[str, str] = {} for qf in (query.measures or []): if not qf.name: continue @@ -1103,8 +1104,8 @@ def _alias_to_short_local(alias: str) -> str: stripped = alias.split(".", 1)[-1] if "." in alias else alias return stripped.replace(".", "__") - _occupied_aliases: Dict[str, str] = {} - _occupied_shorts: Dict[str, str] = {} + _occupied_aliases: dict[str, str] = {} + _occupied_shorts: dict[str, str] = {} for _d in dimensions: _occupied_aliases[_d.alias] = f"dimension '{_d.name}'" _occupied_shorts[_alias_to_short_local(_d.alias)] = ( @@ -1690,7 +1691,7 @@ def _mangled_formula(formula: str) -> str: # operators) flow through unchanged. The construction-time validator # at ``slayer/core/models.py:412`` already rejected DSL constructs. measure_names_set = {m.name for m in measures} - parsed_model_filters: List[ParsedFilter] = [] + parsed_model_filters: list[ParsedFilter] = [] for mf in model.filters: parsed_mf = parse_sql_predicate(mf) for col in parsed_mf.columns: @@ -1733,7 +1734,7 @@ def _mangled_formula(formula: str) -> str: # filters are SQL mode — they don't carry SLayer transforms (rejected # at construction by ``parse_sql_predicate``) and don't go through # ``_preprocess_like`` / ``_preprocess_agg_refs``. - processed_query_filters: List[str] = [] + processed_query_filters: list[str] = [] ft_counter = [0] for f_str in query_filters: rewritten, extra_fields = extract_filter_transforms( @@ -1753,7 +1754,7 @@ def _mangled_formula(formula: str) -> str: # query filters. Used by the windowed-column scan, ``_resolve_joins`` / # ``_collect_needed_paths``, and the ordering of the final # ``EnrichedQuery.filters`` list. - processed_filters_with_mode: List[Tuple[str, str]] = ( + processed_filters_with_mode: list[tuple[str, str]] = ( [(mf, "sql") for mf in model.filters] + [(qf, "dsl") for qf in processed_query_filters] ) @@ -1766,7 +1767,7 @@ def _mangled_formula(formula: str) -> str: # (`rank` / `percent_rank` / `dense_rank` / `ntile`) cover top-N # filtering in pure DSL. Applied to both modes — neither standard SQL # nor SLayer DSL allows window functions in WHERE. - _windowed_columns: Dict[str, str] = { + _windowed_columns: dict[str, str] = { c.name: c.sql for c in model.columns if c.sql and has_window_function(c.sql) } if _windowed_columns: @@ -1800,7 +1801,7 @@ def _mangled_formula(formula: str) -> str: # Names that resolve at the query level (named measures, transforms, # expressions) — pass through as legitimate filter targets even though # they are not Columns / ModelMeasures on the source model. - _query_aliases: Set[str] = set() + _query_aliases: set[str] = set() _query_aliases.update(m.name for m in measures if m.name) _query_aliases.update(t.name for t in enriched_transforms if t.name) _query_aliases.update(e.name for e in enriched_expressions if e.name) @@ -1917,15 +1918,15 @@ def _unpack_dim_resolution(result): async def _maybe_expand( *, - sql: Optional[str], - terminal_model: Optional[SlayerModel], + sql: str | None, + terminal_model: SlayerModel | None, fallback_model: SlayerModel, alias_path: str, resolve_model, named_queries: dict, dialect: str, is_root: bool = True, -) -> Optional[str]: +) -> str | None: """Run the column-SQL expander when we have what we need; otherwise return ``sql`` unchanged. Lets tests that don't supply ``resolve_model`` keep getting the legacy unexpanded behavior — production always supplies @@ -1949,8 +1950,8 @@ async def _maybe_expand( def resolve_via_stage_origin( - *, model: SlayerModel, parts: List[str], -) -> Optional[Column]: + *, model: SlayerModel, parts: list[str], +) -> Column | None: """DEV-1449: Resolve a dotted reference against a virtual stage model produced by ``_query_as_model``. @@ -2008,7 +2009,7 @@ async def _resolve_dotted_dim_with_stage_fallback( model_name_str: str, named_queries: dict, resolve_dimension_via_joins, -) -> "tuple[Optional[Column], Optional[SlayerModel], str]": +) -> "tuple[Column | None, SlayerModel | None, str]": """Resolve a dotted dim / time-dim reference for one query field. Shared by ``_resolve_dimensions`` and ``_resolve_time_dimensions`` @@ -2044,10 +2045,10 @@ async def _resolve_dimensions( resolve_dimension_via_joins, resolve_model=None, dialect: str = "postgres", -) -> List[EnrichedDimension]: +) -> list[EnrichedDimension]: dimensions = [] for dim_ref in query.dimensions or []: - terminal_model: Optional[SlayerModel] = None + terminal_model: SlayerModel | None = None is_local = dim_ref.model is None if is_local: dim_def = model.get_column(dim_ref.name) @@ -2096,10 +2097,10 @@ async def _resolve_time_dimensions( resolve_dimension_via_joins, resolve_model=None, dialect: str = "postgres", -) -> List[EnrichedTimeDimension]: +) -> list[EnrichedTimeDimension]: time_dimensions = [] for td in query.time_dimensions or []: - terminal_model: Optional[SlayerModel] = None + terminal_model: SlayerModel | None = None is_local = td.dimension.model is None if is_local: dim_def = model.get_column(td.dimension.name) @@ -2141,10 +2142,10 @@ async def _resolve_time_dimensions( def _resolve_time_alias( - time_dimensions: List[EnrichedTimeDimension], + time_dimensions: list[EnrichedTimeDimension], query: SlayerQuery, model: SlayerModel, -) -> Optional[str]: +) -> str | None: if len(time_dimensions) == 1: return time_dimensions[0].alias elif len(time_dimensions) > 1: @@ -2162,16 +2163,16 @@ def _resolve_time_alias( def _resolve_last_agg_time( query: SlayerQuery, model: SlayerModel, - dimensions: List[EnrichedDimension], - time_dimensions: List[EnrichedTimeDimension], -) -> Optional[str]: + dimensions: list[EnrichedDimension], + time_dimensions: list[EnrichedTimeDimension], +) -> str | None: if query.main_time_dimension: mtd = query.main_time_dimension if "." not in mtd: mtd = f"{model.name}.{mtd}" return mtd - def _qualified(model_name: str, sql: Optional[str], name: str) -> str: + def _qualified(model_name: str, sql: str | None, name: str) -> str: # Once derived-ref expansion has run, `sql` may already be qualified # (e.g. ``orders.created_at`` instead of bare ``created_at``); don't # double-prefix in that case. @@ -2202,14 +2203,14 @@ def _qualified(model_name: str, sql: Optional[str], name: str) -> str: # --------------------------------------------------------------------------- -def _add_with_prefixes(segments: List[str], paths: Set[Tuple[str, ...]]) -> None: +def _add_with_prefixes(segments: list[str], paths: set[tuple[str, ...]]) -> None: """Add ``segments[:1], segments[:2], …, segments`` to ``paths``.""" for i in range(1, len(segments) + 1): paths.add(tuple(segments[:i])) def _raise_column_cycle( - visited: Tuple[Tuple[str, str], ...], key: Tuple[str, str], + visited: tuple[tuple[str, str], ...], key: tuple[str, str], ) -> None: """Raise a deterministic ``Circular column reference`` error matching the chain format used by ``expand_derived_refs``. @@ -2220,7 +2221,7 @@ def _raise_column_cycle( raise ValueError(f"Circular column reference detected: {chain}") -def _scan_sql_table_refs(*, sql: str, model_name: str, paths: Set[Tuple[str, ...]]) -> None: +def _scan_sql_table_refs(*, sql: str, model_name: str, paths: set[tuple[str, ...]]) -> None: """Regex-fallback scan: pick out ``.`` shapes and add the table prefix paths (skipping references to ``model_name`` itself). """ @@ -2234,9 +2235,9 @@ def _process_node_for_paths( *, node: exp.Column, model: SlayerModel, - paths: Set[Tuple[str, ...]], - visited: Tuple[Tuple[str, str], ...], - dialect: Optional[str] = None, + paths: set[tuple[str, ...]], + visited: tuple[tuple[str, str], ...], + dialect: str | None = None, ) -> None: """Resolve one ``exp.Column`` node into either a recursion into a local derived column or a join-path-prefix add. @@ -2272,9 +2273,9 @@ def _collect_paths_from_local_column_chain( *, model: SlayerModel, col_name: str, - paths: Set[Tuple[str, ...]], - visited: Tuple[Tuple[str, str], ...] = (), - dialect: Optional[str] = None, + paths: set[tuple[str, ...]], + visited: tuple[tuple[str, str], ...] = (), + dialect: str | None = None, ) -> None: """Walk the SQL of a *local* derived column on ``model`` to discover the join paths its expression implies — recursing through references @@ -2318,14 +2319,14 @@ def _collect_paths_from_local_column_chain( def _collect_needed_paths( model: SlayerModel, - dimensions: List[EnrichedDimension], - time_dimensions: List[EnrichedTimeDimension], - measures: List[EnrichedMeasure], + dimensions: list[EnrichedDimension], + time_dimensions: list[EnrichedTimeDimension], + measures: list[EnrichedMeasure], cross_model_measures: list, - processed_filters: List[Tuple[str, str]], - extra_agg_names: Optional[frozenset] = None, - dialect: Optional[str] = None, -) -> Set[Tuple[str, ...]]: + processed_filters: list[tuple[str, str]], + extra_agg_names: frozenset | None = None, + dialect: str | None = None, +) -> set[tuple[str, ...]]: """Extract ordered join-path tuples the query needs (including all prefixes). ``processed_filters`` is a list of ``(filter_text, mode)`` tuples @@ -2334,7 +2335,7 @@ def _collect_needed_paths( matching parser so model filters with arbitrary SQL functions don't trip the DSL allowlist (DEV-1378). """ - paths: Set[Tuple[str, ...]] = set() + paths: set[tuple[str, ...]] = set() for d in dimensions: if d.model_name != model.name: @@ -2380,8 +2381,8 @@ def _scan_filter_column_ref( *, model: SlayerModel, col: str, - paths: Set[Tuple[str, ...]], - dialect: Optional[str] = None, + paths: set[tuple[str, ...]], + dialect: str | None = None, ) -> None: """Route one entry from a parsed filter's column list to the right path-discovery branch. @@ -2404,7 +2405,7 @@ def _scan_filter_column_ref( return if _looks_like_dotted_identifier_ref(col): parts = col.split(".") - expanded: List[str] = [] + expanded: list[str] = [] for part in parts[:-1]: # Model filters convert dots to __; expand both forms. expanded.extend(part.split("__")) @@ -2430,16 +2431,16 @@ def _looks_like_dotted_identifier_ref(value: str) -> bool: async def _resolve_joins( model: SlayerModel, model_name_str: str, - dimensions: List[EnrichedDimension], - time_dimensions: List[EnrichedTimeDimension], - measures: List[EnrichedMeasure], + dimensions: list[EnrichedDimension], + time_dimensions: list[EnrichedTimeDimension], + measures: list[EnrichedMeasure], cross_model_measures: list, - processed_filters: List[Tuple[str, str]], + processed_filters: list[tuple[str, str]], named_queries: dict, resolve_join_target, - extra_agg_names: Optional[frozenset] = None, - dialect: Optional[str] = None, -) -> List[tuple]: + extra_agg_names: frozenset | None = None, + dialect: str | None = None, +) -> list[tuple]: """Resolve only the JOINs the query actually needs by walking the join graph. Instead of relying on baked-in multi-hop joins, this walks each intermediate @@ -2464,8 +2465,8 @@ async def _resolve_joins( # Sort shorter paths first so prefixes are resolved before extensions sorted_paths = sorted(needed_paths, key=len) - resolved_joins: Dict[str, tuple] = {} # alias -> (table_sql, alias, condition) - resolved_models: Dict[str, SlayerModel] = {} # model_name -> SlayerModel + resolved_joins: dict[str, tuple] = {} # alias -> (table_sql, alias, condition) + resolved_models: dict[str, SlayerModel] = {} # model_name -> SlayerModel for path in sorted_paths: alias = "__".join(path) @@ -2533,7 +2534,7 @@ async def _resolve_joins( def _remap_renamed_aliases_in_filter( *, pf: ParsedFilter, - canonical_to_user_name: Dict[str, str], + canonical_to_user_name: dict[str, str], ) -> None: """DEV-1443: rewrite canonical-agg aliases in a parsed query filter to the user-supplied alias when the same node renamed the measure. @@ -2607,7 +2608,7 @@ def _reject_measure_ref_in_filter( *, raw_filter: str, custom_agg_names: frozenset, - named_measures: Dict[str, str], + named_measures: dict[str, str], ) -> None: """DEV-1543: walk a single (substituted) query filter for measure references and raise ``DistinctDimensionValuesError`` on any match. @@ -2666,7 +2667,7 @@ def _reject_measure_ref_in_order_item( *, item: OrderItem, custom_agg_names: frozenset, - named_measures: Dict[str, str], + named_measures: dict[str, str], ) -> None: """DEV-1543: walk a single ``OrderItem`` for measure references. @@ -2722,9 +2723,9 @@ def _reject_measure_ref_in_order_item( def _reject_measure_references_for_raw_rows( *, query: SlayerQuery, - query_filters: List[str], + query_filters: list[str], custom_agg_names: frozenset, - named_measures: Dict[str, str], + named_measures: dict[str, str], ) -> None: """DEV-1543: when ``query.distinct_dimension_values is False``, reject any measure reference in ``query.filters`` or ``query.order``. @@ -2755,9 +2756,9 @@ def _reject_measure_references_for_raw_rows( def extract_filter_transforms( filter_str: str, - counter: Optional[List[int]] = None, - extra_agg_names: Optional[frozenset[str]] = None, - named_measures: Optional[Mapping[str, str]] = None, + counter: list[int] | None = None, + extra_agg_names: frozenset[str] | None = None, + named_measures: Mapping[str, str] | None = None, ) -> tuple: """Extract transform function calls from a filter string. @@ -2811,7 +2812,7 @@ def extract_filter_transforms( except SyntaxError: return filter_str, [] - transforms: List[tuple] = [] + transforms: list[tuple] = [] def _unmangle(s: str) -> str: """Restore colon syntax from placeholders in unparsed formulas.""" @@ -2910,7 +2911,7 @@ async def resolve_filter_columns( *, strict: bool = False, drop_if_unresolved: bool = False, - query_aliases: Optional[Set[str]] = None, + query_aliases: set[str] | None = None, ) -> list: """Resolve filter column references through model dimensions/measures. @@ -3243,9 +3244,9 @@ def _classify_one_filter( def classify_filters( filters: list, measure_names: set, - computed_names: Optional[set] = None, - groupby_names: Optional[set] = None, - windowed_measure_names: Optional[set] = None, + computed_names: set | None = None, + groupby_names: set | None = None, + windowed_measure_names: set | None = None, ) -> list: """Classify filters as WHERE, HAVING, or post-filter. diff --git a/slayer/engine/ingestion.py b/slayer/engine/ingestion.py index bcc259d9..fa5dd485 100644 --- a/slayer/engine/ingestion.py +++ b/slayer/engine/ingestion.py @@ -11,7 +11,7 @@ import logging import sys from collections import defaultdict, deque -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, TextIO, Tuple +from typing import TYPE_CHECKING, Any, TextIO import sqlalchemy as sa import sqlalchemy.dialects.mssql as _sqla_mssql @@ -41,7 +41,7 @@ # Module-level dedup set for unrecognized SA type warnings (see # _sa_type_to_data_type). Keyed by upper-cased class name. -_logged_unmapped_sa_types: Set[str] = set() +_logged_unmapped_sa_types: set[str] = set() # Map SQLAlchemy types to SLayer DataTypes. # DEV-1361: integer family → INT, floating family → DOUBLE, NUMERIC/DECIMAL @@ -268,9 +268,9 @@ class RollupGraphError(Exception): def _get_fk_relationships( inspector: sa.engine.Inspector, table_name: str, - schema: Optional[str], - table_set: Set[str], -) -> List[tuple]: + schema: str | None, + table_set: set[str], +) -> list[tuple]: """Get FK relationships for a table, filtered to tables in table_set. Returns list of (source_column, target_table, target_column). @@ -290,12 +290,12 @@ def _get_fk_relationships( def _build_fk_graph( inspector: sa.engine.Inspector, - table_names: List[str], - schema: Optional[str], -) -> Dict[str, Set[str]]: + table_names: list[str], + schema: str | None, +) -> dict[str, set[str]]: """Build directed graph: graph[table] = set of tables it references via FK.""" table_set = set(table_names) - graph: Dict[str, Set[str]] = defaultdict(set) + graph: dict[str, set[str]] = defaultdict(set) for table_name in table_names: for _, ref_table, _ in _get_fk_relationships( inspector=inspector, @@ -307,12 +307,12 @@ def _build_fk_graph( return dict(graph) -def _check_acyclic(graph: Dict[str, Set[str]]) -> None: +def _check_acyclic(graph: dict[str, set[str]]) -> None: """Check that FK graph is a DAG. Raises RollupGraphError if cycles found.""" - visited: Set[str] = set() - rec_stack: Set[str] = set() + visited: set[str] = set() + rec_stack: set[str] = set() - def dfs(node: str, path: List[str]) -> None: + def dfs(node: str, path: list[str]) -> None: visited.add(node) rec_stack.add(node) path.append(node) @@ -326,7 +326,7 @@ def dfs(node: str, path: List[str]) -> None: path.pop() rec_stack.remove(node) - all_nodes: Set[str] = set(graph.keys()) + all_nodes: set[str] = set(graph.keys()) for neighbors in graph.values(): all_nodes.update(neighbors) for node in all_nodes: @@ -334,9 +334,9 @@ def dfs(node: str, path: List[str]) -> None: dfs(node, []) -def _compute_transitive_closure(graph: Dict[str, Set[str]], source: str) -> Set[str]: +def _compute_transitive_closure(graph: dict[str, set[str]], source: str) -> set[str]: """BFS to find all tables transitively reachable from source (excluding source).""" - reachable: Set[str] = set() + reachable: set[str] = set() queue = deque([source]) visited = {source} while queue: @@ -357,10 +357,10 @@ def _compute_transitive_closure(graph: Dict[str, Set[str]], source: str) -> Set[ def _generate_joins( inspector: sa.engine.Inspector, source_table: str, - referenced_tables: Set[str], - schema: Optional[str], - table_set: Set[str], -) -> List[ModelJoin]: + referenced_tables: set[str], + schema: str | None, + table_set: set[str], +) -> list[ModelJoin]: """Generate direct ModelJoin objects from the source table's own FK relationships. Only emits joins for FKs defined on ``source_table`` itself — multi-hop @@ -375,7 +375,7 @@ def _generate_joins( ) joins = [] - seen_signatures: Set[Tuple[str, str, str]] = set() + seen_signatures: set[tuple[str, str, str]] = set() for src_col, ref_table, tgt_col in fk_rels: if ref_table not in referenced_tables: continue @@ -418,8 +418,8 @@ def _parse_info_schema_is_float(data_type_str: str) -> bool: def _get_columns_fallback( sa_engine: sa.Engine, table_name: str, - schema: Optional[str], -) -> List[Dict]: + schema: str | None, +) -> list[dict]: """Get columns via INFORMATION_SCHEMA when Inspector.get_columns() fails.""" if schema: sql = ( @@ -466,8 +466,8 @@ def _get_columns_fallback( def _get_pk_constraint_fallback( sa_engine: sa.Engine, table_name: str, - schema: Optional[str], -) -> Dict: + schema: str | None, +) -> dict: """Get PK constraint via INFORMATION_SCHEMA when Inspector.get_pk_constraint() fails.""" if schema: sql = ( @@ -501,8 +501,8 @@ def _safe_get_columns( inspector: sa.engine.Inspector, sa_engine: sa.Engine, table_name: str, - schema: Optional[str], -) -> List[Dict]: + schema: str | None, +) -> list[dict]: """Get columns, falling back to INFORMATION_SCHEMA on failure.""" try: return inspector.get_columns(table_name, schema=schema) @@ -514,8 +514,8 @@ def _safe_get_pk_constraint( inspector: sa.engine.Inspector, sa_engine: sa.Engine, table_name: str, - schema: Optional[str], -) -> Dict: + schema: str | None, +) -> dict: """Get PK constraint, falling back to INFORMATION_SCHEMA on failure. SQLite has no information_schema views; its stock inspector reads @@ -541,12 +541,12 @@ def _introspect_query_columns_via_inspector( sa_engine: sa.Engine, inspector: sa.engine.Inspector, table_name: str, - schema: Optional[str], - rollup_sql: Optional[str], - referenced_tables: Set[str], - fk_columns_by_table: Dict[str, Set[str]], - joins: Optional[List[ModelJoin]] = None, -) -> List[tuple]: + schema: str | None, + rollup_sql: str | None, + referenced_tables: set[str], + fk_columns_by_table: dict[str, set[str]], + joins: list[ModelJoin] | None = None, +) -> list[tuple]: """Introspect columns from a rollup query or plain table. Returns list of (column_name, DataType, is_primary_key, is_float) tuples. @@ -574,7 +574,7 @@ def _introspect_query_columns_via_inspector( # Build list of (ref_table, dotted_path) from joins — supports diamond joins # where the same table appears via multiple paths - table_path_pairs: List[tuple] = [] + table_path_pairs: list[tuple] = [] if joins: for mj in joins: if mj.join_pairs and "." in mj.join_pairs[0][0]: @@ -619,10 +619,10 @@ def _introspect_query_columns_via_inspector( def _columns_to_model( name: str, - columns: List[tuple], + columns: list[tuple], data_source: str, - sql_table: Optional[str] = None, - joins: Optional[List[ModelJoin]] = None, + sql_table: str | None = None, + joins: list[ModelJoin] | None = None, ) -> SlayerModel: """Generate a SlayerModel from introspected (column_name, DataType, is_pk, is_float) tuples. @@ -630,7 +630,7 @@ def _columns_to_model( used as is decided per query. This function emits one Column per non-joined column, with format inferred from the column's data type. """ - cols: List[Column] = [] + cols: list[Column] = [] _INT_FORMAT = NumberFormat(type=NumberFormatType.INTEGER) _FLOAT_FORMAT = NumberFormat(type=NumberFormatType.FLOAT) @@ -675,8 +675,8 @@ def _sqlite_probe_integer_columns( *, sa_engine: sa.Engine, sql_table: str, - columns: List[tuple], -) -> List[tuple]: + columns: list[tuple], +) -> list[tuple]: """DEV-1538: per-column SQLite affinity probe. Walks the tuples ``(col_name, DataType, is_pk, is_float)`` produced by @@ -703,7 +703,7 @@ def _sqlite_probe_integer_columns( from slayer.sql.sqlite_introspect import probe_sqlite_integer_column schema, table = _parse_qualified_sql_table(sql_table) - out: List[tuple] = [] + out: list[tuple] = [] with sa_engine.connect() as conn: for col_name, data_type, is_pk, is_float in columns: if data_type is not DataType.INT or "." in col_name: @@ -736,7 +736,7 @@ def _sqlite_probe_integer_columns( return out -def _parse_qualified_sql_table(sql_table: str) -> Tuple[Optional[str], str]: +def _parse_qualified_sql_table(sql_table: str) -> tuple[str | None, str]: """Split ``"schema.table"`` into ``(schema, table)`` or ``(None, table)``. Only splits on a single dot — table/schema names containing dots are @@ -754,9 +754,9 @@ def introspect_table_to_model( sa_engine: sa.Engine, inspector: sa.engine.Inspector, table_name: str, - schema: Optional[str], + schema: str | None, data_source: str, - model_name: Optional[str] = None, + model_name: str | None = None, ) -> SlayerModel: """Introspect a single table (no FK rollup) and return a SlayerModel. @@ -793,10 +793,10 @@ def introspect_table_to_model( def ingest_datasource( datasource: DatasourceConfig, - include_tables: Optional[List[str]] = None, - exclude_tables: Optional[List[str]] = None, - schema: Optional[str] = None, -) -> List[SlayerModel]: + include_tables: list[str] | None = None, + exclude_tables: list[str] | None = None, + schema: str | None = None, +) -> list[SlayerModel]: from slayer.sql import engine_factory sa_engine = engine_factory.get_engine(datasource.resolve_env_vars()) inspector = sa.inspect(sa_engine) @@ -819,7 +819,7 @@ def ingest_datasource( has_cycles = True # Collect FK columns per table (for excluding from rollup) - fk_columns_by_table: Dict[str, Set[str]] = defaultdict(set) + fk_columns_by_table: dict[str, set[str]] = defaultdict(set) for table_name in table_names: fks = inspector.get_foreign_keys(table_name, schema=schema) for fk in fks: @@ -902,18 +902,18 @@ def ingest_datasource( # --------------------------------------------------------------------------- -def _existing_join_signatures(model: SlayerModel) -> Set[Tuple[str, Tuple[Tuple[str, str], ...]]]: +def _existing_join_signatures(model: SlayerModel) -> set[tuple[str, tuple[tuple[str, str], ...]]]: """Return the set of (target_model, sorted join_pair tuples) signatures for joins already on ``model``. Used to detect new joins. """ - out: Set[Tuple[str, Tuple[Tuple[str, str], ...]]] = set() + out: set[tuple[str, tuple[tuple[str, str], ...]]] = set() for j in model.joins: sig_pairs = tuple(sorted((p[0], p[1]) for p in j.join_pairs)) out.add((j.target_model, sig_pairs)) return out -def _is_auto_default_integer_format(fmt: Optional[NumberFormat]) -> bool: +def _is_auto_default_integer_format(fmt: NumberFormat | None) -> bool: """Return True when ``fmt`` looks like the auto-ingested ``NumberFormat (type=INTEGER)`` default (no custom precision / symbol set). Used by DEV-1538's widening path to decide whether to flip the format alongside @@ -926,7 +926,7 @@ def _is_auto_default_integer_format(fmt: Optional[NumberFormat]) -> bool: return fmt.precision is None and fmt.symbol is None -def _format_for_widened_type(verdict: DataType) -> Optional[NumberFormat]: +def _format_for_widened_type(verdict: DataType) -> NumberFormat | None: """Return the auto-default format for a probed widening verdict.""" if verdict is DataType.DOUBLE: return NumberFormat(type=NumberFormatType.FLOAT) @@ -936,10 +936,10 @@ def _format_for_widened_type(verdict: DataType) -> Optional[NumberFormat]: def _merge_persisted_column_with_probe( *, persisted_col: Column, - fresh_col: Optional[Column], + fresh_col: Column | None, model_name: str, sqlite_widen_enabled: bool, -) -> Tuple[Column, bool]: +) -> tuple[Column, bool]: """DEV-1538: decide whether a persisted column should be widened based on a freshly-probed type, and return ``(merged_column, did_widen)``. @@ -956,7 +956,7 @@ def _merge_persisted_column_with_probe( ): return persisted_col, False - updates: Dict[str, Any] = {"type": fresh_col.type} + updates: dict[str, Any] = {"type": fresh_col.type} if _is_auto_default_integer_format(persisted_col.format): updates["format"] = _format_for_widened_type(fresh_col.type) else: @@ -972,14 +972,14 @@ def _merge_persisted_column_with_probe( def _merge_joins_strict( persisted: SlayerModel, fresh: SlayerModel, -) -> Tuple[List[ModelJoin], List[str]]: +) -> tuple[list[ModelJoin], list[str]]: """Append joins whose signature isn't already present. Raises on the duplicate-target / different-pairs conflict so callers don't end up with two joins pointing at the same target_model.""" existing_join_sigs = _existing_join_signatures(persisted) existing_join_targets = {j.target_model for j in persisted.joins} - new_joins: List[ModelJoin] = list(persisted.joins) - new_join_targets: List[str] = [] + new_joins: list[ModelJoin] = list(persisted.joins) + new_join_targets: list[str] = [] for j in fresh.joins: sig = (j.target_model, tuple(sorted((p[0], p[1]) for p in j.join_pairs))) if sig in existing_join_sigs: @@ -1003,7 +1003,7 @@ def _additive_merge_existing( persisted: SlayerModel, fresh: SlayerModel, sqlite_widen_enabled: bool = False, -) -> Tuple[SlayerModel, List[str], List[str], List[str]]: +) -> tuple[SlayerModel, list[str], list[str], list[str]]: """Merge a freshly-ingested ``fresh`` model into ``persisted`` additively. Returns ``(merged, new_column_names, new_join_target_names, @@ -1024,11 +1024,11 @@ def _additive_merge_existing( appended from ``fresh.columns``. * Joins with new ``(target_model, join_pairs)`` signatures are appended. """ - existing_by_name: Dict[str, Column] = {c.name: c for c in persisted.columns} - fresh_by_name: Dict[str, Column] = {c.name: c for c in fresh.columns} + existing_by_name: dict[str, Column] = {c.name: c for c in persisted.columns} + fresh_by_name: dict[str, Column] = {c.name: c for c in fresh.columns} - widened_column_names: List[str] = [] - merged_columns: List[Column] = [] + widened_column_names: list[str] = [] + merged_columns: list[Column] = [] for persisted_col in persisted.columns: merged_col, did_widen = _merge_persisted_column_with_probe( persisted_col=persisted_col, @@ -1040,7 +1040,7 @@ def _additive_merge_existing( if did_widen: widened_column_names.append(persisted_col.name) - new_column_names: List[str] = [] + new_column_names: list[str] = [] for fresh_col in fresh.columns: if fresh_col.name in existing_by_name: continue @@ -1111,8 +1111,8 @@ async def _scoped_models_for_validation( *, storage: StorageBackend, datasource: DatasourceConfig, - in_scope_table_names: Set[str], -) -> List[SlayerModel]: + in_scope_table_names: set[str], +) -> list[SlayerModel]: """Build the list of persisted models to feed to ``validate_datasource``. sql_table-mode models are included only when their live table is in @@ -1122,7 +1122,7 @@ async def _scoped_models_for_validation( """ identities = await storage._list_all_model_identities() ds_model_names = [n for d, n in identities if d == datasource.name] - scoped: List[SlayerModel] = [] + scoped: list[SlayerModel] = [] for name in ds_model_names: m = await storage.get_model(name, data_source=datasource.name) if m is None: @@ -1139,9 +1139,9 @@ async def ingest_datasource_idempotent( *, datasource: DatasourceConfig, storage: StorageBackend, - include_tables: Optional[List[str]] = None, - exclude_tables: Optional[List[str]] = None, - schema: Optional[str] = None, + include_tables: list[str] | None = None, + exclude_tables: list[str] | None = None, + schema: str | None = None, ): """Idempotent re-ingestion (DEV-1356). @@ -1165,8 +1165,8 @@ async def ingest_datasource_idempotent( validate_datasource, ) - additions: List[ModelAddition] = [] - errors: List[IngestionError] = [] + additions: list[ModelAddition] = [] + errors: list[IngestionError] = [] # ``ingest_datasource`` is sync (it drives SQLAlchemy ``Inspector``). # Offload to a thread so a slow / large datasource doesn't block the @@ -1179,7 +1179,7 @@ async def ingest_datasource_idempotent( schema=schema, ) fresh_by_name = {m.name: m for m in fresh_models} - in_scope_table_names: Set[str] = set(fresh_by_name.keys()) + in_scope_table_names: set[str] = set(fresh_by_name.keys()) for table_name, fresh in fresh_by_name.items(): try: @@ -1296,7 +1296,7 @@ def _friendly_db_error(exc: Exception) -> str: def _print_ingest_addition( - addition, *, file: Optional[TextIO] = None + addition, *, file: TextIO | None = None ) -> None: out = file if file is not None else sys.stdout if addition.created: @@ -1319,7 +1319,7 @@ def _print_ingest_addition( def _print_ingest_drift_and_errors( - result, *, file: Optional[TextIO] = None + result, *, file: TextIO | None = None ) -> None: out = file if file is not None else sys.stdout if result.to_delete: @@ -1354,15 +1354,15 @@ class StartupIngestSummary(BaseModel): ``EditModelDelete | WholeModelDelete``. """ - succeeded: List[str] = Field(default_factory=list) - failures: List[StartupIngestFailure] = Field(default_factory=list) - drift_pending: List[Any] = Field(default_factory=list) + succeeded: list[str] = Field(default_factory=list) + failures: list[StartupIngestFailure] = Field(default_factory=list) + drift_pending: list[Any] = Field(default_factory=list) async def ingest_all_datasources_idempotent( *, storage: StorageBackend, - stream: Optional[TextIO] = None, + stream: TextIO | None = None, ) -> StartupIngestSummary: """Run idempotent auto-ingestion across every configured datasource. @@ -1436,7 +1436,7 @@ async def _refresh_models_for_datasource( datasource_name: str, storage: StorageBackend, search: "SearchService", -) -> Tuple[List[Tuple[str, str]], List[SlayerModel]]: +) -> tuple[list[tuple[str, str]], list[SlayerModel]]: """Refresh embeddings for every visible model in the datasource. Returns ``(warnings, models_in_ds)``. Each warning is tagged with @@ -1444,8 +1444,8 @@ async def _refresh_models_for_datasource( right ``IngestionError.model_name``. ``models_in_ds`` is forwarded to the datasource-doc refresh that follows. """ - warnings: List[Tuple[str, str]] = [] - models_in_ds: List[SlayerModel] = [] + warnings: list[tuple[str, str]] = [] + models_in_ds: list[SlayerModel] = [] try: identities = await storage._list_all_model_identities() except Exception as exc: # noqa: BLE001 — defensive @@ -1474,10 +1474,10 @@ async def _refresh_models_for_datasource( async def _refresh_datasource_doc( *, datasource_name: str, - models: List[SlayerModel], + models: list[SlayerModel], search: "SearchService", storage: StorageBackend, -) -> List[Tuple[str, str]]: +) -> list[tuple[str, str]]: """Refresh the datasource doc embedding. Warnings are tagged with an empty ``model_name`` since the doc has no specific entity name. @@ -1497,7 +1497,7 @@ async def _refresh_datasource_doc( async def _entity_ref_exists( *, entity: str, storage: StorageBackend, -) -> Optional[bool]: +) -> bool | None: """DEV-1428 defense-in-depth cleanup probe. Returns: * ``True`` when the canonical ref still resolves. @@ -1548,7 +1548,7 @@ async def _refresh_memories_for_datasource( # NOSONAR(S3776) — straight-line datasource_name: str, storage: StorageBackend, search: "SearchService", -) -> List[Tuple[str, str]]: +) -> list[tuple[str, str]]: """Refresh embeddings for every memory whose canonical entities are rooted at this datasource. Each warning is tagged with ``memory:`` so a startup log inspection can distinguish memory @@ -1568,7 +1568,7 @@ async def _refresh_memories_for_datasource( # NOSONAR(S3776) — straight-line memories = await storage.list_memories() except Exception as exc: # noqa: BLE001 — defensive return [("", f"{datasource_name} (memories): {exc}")] - warnings: List[Tuple[str, str]] = [] + warnings: list[tuple[str, str]] = [] for memory in memories: rooted_at_ds = any( canonical_id_rooted_at(e, datasource_name) @@ -1595,7 +1595,7 @@ async def _refresh_memories_for_datasource( # NOSONAR(S3776) — straight-line warnings.append((tag, w)) # DEV-1428 cleanup pass: drop refs that resolve to False # (definitive not-found); keep refs that raise (transient). - cleaned: List[str] = [] + cleaned: list[str] = [] changed = False for entity in memory.entities: exists = await _entity_ref_exists( @@ -1626,7 +1626,7 @@ async def _refresh_memories_for_datasource( # NOSONAR(S3776) — straight-line async def _refresh_datasource_embeddings( *, datasource_name: str, storage: StorageBackend, -) -> List[Tuple[str, str]]: +) -> list[tuple[str, str]]: """Refresh persisted embeddings for everything reachable from this datasource: every visible model + its visible children, the datasource doc itself, and every memory whose canonical entities diff --git a/slayer/engine/profiling.py b/slayer/engine/profiling.py index 9fb73c97..5ba343b4 100644 --- a/slayer/engine/profiling.py +++ b/slayer/engine/profiling.py @@ -57,7 +57,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple +from typing import Any, NamedTuple from slayer.core.enums import DataType from slayer.core.models import Column, SlayerModel @@ -90,9 +90,9 @@ class ColumnSample(NamedTuple): ``None`` for numeric/temporal columns. """ - sampled: Optional[str] - sampled_values: Optional[List[str]] - distinct_count: Optional[int] + sampled: str | None + sampled_values: list[str] | None + distinct_count: int | None # --------------------------------------------------------------------------- @@ -111,10 +111,10 @@ class _DimProfileEntry(NamedTuple): name: str type_str: str - distinct_count: Optional[int] - values: Optional[List[Any]] - min_value: Optional[Any] - max_value: Optional[Any] + distinct_count: int | None + values: list[Any] | None + min_value: Any | None + max_value: Any | None def _format_dim_profile_value(entry: _DimProfileEntry) -> str: @@ -151,7 +151,7 @@ async def _profile_categorical_column( column: Column, engine: SlayerQueryEngine, max_values: int, -) -> Optional[_DimProfileEntry]: +) -> _DimProfileEntry | None: """Profile one string/boolean column. DEV-1480: orders by per-value count desc with alphabetical tie-break in @@ -181,7 +181,7 @@ async def _profile_categorical_column( value_key = f"{model.name}.{column.name}" # Filter NULL values out — they map to ``col IS NULL`` predicates, not # to literal-equality use cases the validator cares about. - raw_pairs: List[Tuple[Any, Any]] = [] + raw_pairs: list[tuple[Any, Any]] = [] count_key = f"{model.name}._count" for row in r.data: v = row.get(value_key) @@ -193,7 +193,7 @@ async def _profile_categorical_column( # equally-ranked rows in arbitrary order. NB: this only re-orders what # we received — the LIMIT cutoff is the SQL's responsibility. raw_pairs.sort(key=lambda p: (-(p[1] or 0), str(p[0]))) - values: List[str] = [str(v) for v, _ in raw_pairs] + values: list[str] = [str(v) for v, _ in raw_pairs] overflow = len(values) > max_values return _DimProfileEntry( name=column.name, @@ -208,9 +208,9 @@ async def _profile_categorical_column( async def _profile_numeric_temporal_columns( *, model: SlayerModel, - columns: List[Column], + columns: list[Column], engine: SlayerQueryEngine, -) -> Dict[str, _DimProfileEntry]: +) -> dict[str, _DimProfileEntry]: """Profile every numeric/temporal column in a single batched min/max query.""" if not columns: return {} @@ -219,11 +219,11 @@ async def _profile_numeric_temporal_columns( "type": str(c.type)} for c in columns ] - measures_payload: List[Dict[str, str]] = [] + measures_payload: list[dict[str, str]] = [] for c in columns: measures_payload.append({"formula": f"_slayer_range_{c.name}:min"}) measures_payload.append({"formula": f"_slayer_range_{c.name}:max"}) - row: Dict[str, Any] = {} + row: dict[str, Any] = {} try: q = SlayerQuery.model_validate({ "source_model": {"source_name": model.name, "columns": ext_columns}, @@ -234,7 +234,7 @@ async def _profile_numeric_temporal_columns( row = r.data[0] except Exception: row = {} - out: Dict[str, _DimProfileEntry] = {} + out: dict[str, _DimProfileEntry] = {} for c in columns: mn = row.get(f"{model.name}._slayer_range_{c.name}_min") mx = row.get(f"{model.name}._slayer_range_{c.name}_max") @@ -257,8 +257,8 @@ async def _collect_dim_profile( engine: SlayerQueryEngine, max_values: int = _MAX_CATEGORICAL_VALUES, max_dims: int = 10, - only_columns: Optional[Set[str]] = None, -) -> List[_DimProfileEntry]: + only_columns: set[str] | None = None, +) -> list[_DimProfileEntry]: """Produce one profile entry per eligible column (non-hidden, non-pk). - string/boolean columns: distinct values (or overflow marker) via one @@ -288,7 +288,7 @@ async def _collect_dim_profile( if c.type in (DataType.INT, DataType.DOUBLE, DataType.DATE, DataType.TIMESTAMP) ] - entries: Dict[str, _DimProfileEntry] = {} + entries: dict[str, _DimProfileEntry] = {} for c in categorical: entry = await _profile_categorical_column( model=model, column=c, engine=engine, max_values=max_values, @@ -353,7 +353,7 @@ async def _count_distinct_via_model_extension( model: SlayerModel, column: Column, engine: SlayerQueryEngine, -) -> Optional[int]: +) -> int | None: """Fire a secondary ``count_distinct`` query via a transient ``ModelExtension`` column. @@ -399,7 +399,7 @@ async def _profile_categorical_with_total( model: SlayerModel, column: Column, engine: SlayerQueryEngine, -) -> Optional[ColumnSample]: +) -> ColumnSample | None: """DEV-1480 categorical profile: top-50 by frequency + true total on overflow. @@ -426,14 +426,14 @@ async def _profile_categorical_with_total( return None value_key = f"{model.name}.{column.name}" count_key = f"{model.name}._count" - raw_pairs: List[Tuple[Any, Any]] = [] + raw_pairs: list[tuple[Any, Any]] = [] for row in r.data: v = row.get(value_key) if v is None: continue raw_pairs.append((v, row.get(count_key))) raw_pairs.sort(key=lambda p: (-(p[1] or 0), str(p[0]))) - values: List[str] = [str(v) for v, _ in raw_pairs] + values: list[str] = [str(v) for v, _ in raw_pairs] overflow = len(values) > _MAX_CATEGORICAL_VALUES if not overflow: text = ", ".join(values[:_TEXT_SAMPLE_CAP]) @@ -473,7 +473,7 @@ async def profile_column( model: SlayerModel, column: Column, engine: SlayerQueryEngine, -) -> Optional[ColumnSample]: +) -> ColumnSample | None: """Return the :class:`ColumnSample` for ``column`` on ``model``. Returns ``None`` for primary-key / hidden columns and when the @@ -513,14 +513,14 @@ async def _refresh_one_column( column: Column, engine: SlayerQueryEngine, storage: StorageBackend, -) -> List[str]: +) -> list[str]: """Profile + persist a single column. Best-effort — returns the list of error strings produced (empty on full success). Extracted from ``refresh_table_backed_model_sampled`` to keep that function's cognitive complexity low. """ - errors: List[str] = [] - sample: Optional[ColumnSample] = None + errors: list[str] = [] + sample: ColumnSample | None = None try: sample = await profile_column(model=model, column=column, engine=engine) except Exception as exc: # NOSONAR(S112) — best-effort: see module docstring @@ -550,8 +550,8 @@ async def refresh_table_backed_model_sampled( model: SlayerModel, engine: SlayerQueryEngine, storage: StorageBackend, - only_columns: Optional[Set[str]] = None, -) -> List[str]: + only_columns: set[str] | None = None, +) -> list[str]: """Refresh ``Column.sampled``, ``Column.sampled_values``, and ``Column.distinct_count`` for each eligible column on ``model``. @@ -562,7 +562,7 @@ async def refresh_table_backed_model_sampled( """ if not _is_table_backed(model): return [] - errors: List[str] = [] + errors: list[str] = [] for column in model.columns: if column.hidden or column.primary_key: continue @@ -579,10 +579,10 @@ async def refresh_all_table_backed_sampled( engine: SlayerQueryEngine, storage: StorageBackend, data_source: str, -) -> List[str]: +) -> list[str]: """Refresh ``Column.sampled`` for every table-backed model in ``data_source``. Best-effort across all models.""" - errors: List[str] = [] + errors: list[str] = [] identities = await storage._list_all_model_identities() for ds, name in identities: if ds != data_source: @@ -604,9 +604,9 @@ async def handle_edit_refresh( storage: StorageBackend, data_source: str, model_name: str, - changed_columns: Set[str], + changed_columns: set[str], model_level_change: bool, -) -> List[str]: +) -> list[str]: """Refresh entry point for ``edit_model``. * ``model_level_change=True`` → refresh every non-hidden column on diff --git a/slayer/engine/query_engine.py b/slayer/engine/query_engine.py index 991921b3..8466a712 100644 --- a/slayer/engine/query_engine.py +++ b/slayer/engine/query_engine.py @@ -6,7 +6,7 @@ import decimal import logging from contextvars import ContextVar -from typing import Any, Dict, List, Optional, Tuple +from typing import Any from pydantic import BaseModel, Field as PydanticField, model_validator from sqlglot import exp @@ -44,7 +44,7 @@ logger = logging.getLogger(__name__) -def _sql_client_cache_key(datasource: DatasourceConfig) -> Tuple[str, str]: +def _sql_client_cache_key(datasource: DatasourceConfig) -> tuple[str, str]: """Cache key for ``SlayerQueryEngine._sql_clients``. Mirrors ``engine_factory``'s cache key so two datasources differing @@ -61,7 +61,7 @@ def _sql_client_cache_key(datasource: DatasourceConfig) -> Tuple[str, str]: # engine don't see each other's in-flight state — each asyncio task gets its # own copy of the context. The default=None + lazy-init pattern below means # only tasks that actually hit a query-backed join target allocate a set. -_join_target_resolving_var: ContextVar[Optional[set]] = ContextVar( +_join_target_resolving_var: ContextVar[set | None] = ContextVar( "_join_target_resolving", default=None ) @@ -72,7 +72,7 @@ def _sql_client_cache_key(datasource: DatasourceConfig) -> Tuple[str, str]: # _resolve_model_inner to differentiate forward/self refs from genuine # misspellings, so the user gets a clear error instead of "Model 'X' not found". # Each entry maps a forbidden target name to the stage that tried to reach it. -_forbidden_sibling_refs_var: ContextVar[Optional[Dict[str, str]]] = ContextVar( +_forbidden_sibling_refs_var: ContextVar[dict[str, str] | None] = ContextVar( "_forbidden_sibling_refs", default=None ) @@ -97,10 +97,10 @@ def __init__(self, hop_name: str) -> None: def _merge_query_variables( *, - outer: Optional[Dict[str, Any]], - stage: Optional[Dict[str, Any]], - runtime: Optional[Dict[str, Any]], -) -> Dict[str, Any]: + outer: dict[str, Any] | None, + stage: dict[str, Any] | None, + runtime: dict[str, Any] | None, +) -> dict[str, Any]: """Merge variable layers per spec precedence: ``runtime > stage > outer``. Model-level defaults are folded into ``outer`` by the caller before @@ -110,8 +110,8 @@ def _merge_query_variables( def _apply_placeholder_fill( - query: SlayerQuery, effective: Dict[str, Any] -) -> Dict[str, Any]: + query: SlayerQuery, effective: dict[str, Any] +) -> dict[str, Any]: """Add ``{var: '0'}`` for any unresolved ``{var}`` placeholder in ``query.filters`` so save-time dry-run SQL generation can proceed even when a runtime variable has no default. @@ -148,17 +148,17 @@ def _build_explain_sql(dialect: str, sql: str) -> str: class FieldMetadata(BaseModel): """Metadata for a single field in the query response.""" - label: Optional[str] = None - format: Optional[NumberFormat] = None + label: str | None = None + format: NumberFormat | None = None class ResponseAttributes(BaseModel): """Field metadata for a query response, split by type.""" - dimensions: Dict[str, FieldMetadata] = PydanticField(default_factory=dict) - measures: Dict[str, FieldMetadata] = PydanticField(default_factory=dict) + dimensions: dict[str, FieldMetadata] = PydanticField(default_factory=dict) + measures: dict[str, FieldMetadata] = PydanticField(default_factory=dict) - def get(self, column: str) -> Optional[FieldMetadata]: + def get(self, column: str) -> FieldMetadata | None: """Look up metadata for a column across both dicts.""" return self.dimensions.get(column) or self.measures.get(column) @@ -166,9 +166,9 @@ def get(self, column: str) -> Optional[FieldMetadata]: class SlayerResponse(BaseModel): """Response from a SLayer query.""" - data: List[Dict[str, Any]] - columns: List[str] = PydanticField(default_factory=list) - sql: Optional[str] = None + data: list[dict[str, Any]] + columns: list[str] = PydanticField(default_factory=list) + sql: str | None = None attributes: ResponseAttributes = PydanticField(default_factory=ResponseAttributes) @model_validator(mode="after") @@ -207,7 +207,7 @@ def _infer_aggregated_format( model: SlayerModel, measure_name: str, aggregation: str, -) -> Optional[NumberFormat]: +) -> NumberFormat | None: """Infer NumberFormat for an aggregated measure based on aggregation type and source measure format. Rules: @@ -247,7 +247,7 @@ def __init__(self, storage: StorageBackend): # ``engine_factory``'s cache so Snowflake datasources sharing a # connection_name but differing in warehouse/role/database/schema # get distinct clients (DEV-1551). - self._sql_clients: Dict[Tuple[str, str], SlayerSQLClient] = {} + self._sql_clients: dict[tuple[str, str], SlayerSQLClient] = {} def _get_join_target_resolving(self) -> set: """Return the per-task in-flight join-target name set, allocating one @@ -261,8 +261,8 @@ def _get_join_target_resolving(self) -> set: @staticmethod def _scope_named_queries_to_prior( - named_queries: Dict[str, "SlayerQuery"], stage_name: Optional[str] - ) -> Dict[str, "SlayerQuery"]: + named_queries: dict[str, "SlayerQuery"], stage_name: str | None + ) -> dict[str, "SlayerQuery"]: """Slice an insertion-ordered named-queries dict to entries that come strictly before ``stage_name``. @@ -281,7 +281,7 @@ def _scope_named_queries_to_prior( """ if not stage_name or stage_name not in named_queries: return named_queries - out: Dict[str, "SlayerQuery"] = {} + out: dict[str, "SlayerQuery"] = {} for k, v in named_queries.items(): if k == stage_name: return out @@ -329,13 +329,13 @@ def _extract_sibling_refs(query: "SlayerQuery", against: set) -> set: @staticmethod def _index_query_list_by_name( - rest: List["SlayerQuery"], root: "SlayerQuery", - ) -> Dict[str, "SlayerQuery"]: + rest: list["SlayerQuery"], root: "SlayerQuery", + ) -> dict[str, "SlayerQuery"]: """Build ``{name: query}`` for non-final entries, validating that every non-final entry has a unique name and that the root's name (if any) doesn't collide. """ - rest_by_name: Dict[str, "SlayerQuery"] = {} + rest_by_name: dict[str, "SlayerQuery"] = {} for q in rest: if not q.name: raise ValueError( @@ -355,8 +355,8 @@ def _index_query_list_by_name( @classmethod def _validate_query_list_invariants( cls, - queries: List["SlayerQuery"], - rest: List["SlayerQuery"], + queries: list["SlayerQuery"], + rest: list["SlayerQuery"], root: "SlayerQuery", sibling_names: set, ) -> None: @@ -385,7 +385,7 @@ def _validate_query_list_invariants( @classmethod def _build_dependency_graph( cls, - rest_by_name: Dict[str, "SlayerQuery"], + rest_by_name: dict[str, "SlayerQuery"], sibling_names: set, ) -> tuple: """Build the (in_degree, dependents) adjacency for Kahn's. @@ -394,8 +394,8 @@ def _build_dependency_graph( the count of siblings ``X`` depends on; ``dependents[X]`` is the list of siblings that depend on ``X``. """ - in_degree: Dict[str, int] = dict.fromkeys(rest_by_name, 0) - dependents: Dict[str, List[str]] = {name: [] for name in rest_by_name} + in_degree: dict[str, int] = dict.fromkeys(rest_by_name, 0) + dependents: dict[str, list[str]] = {name: [] for name in rest_by_name} for name, q in rest_by_name.items(): for prereq in cls._extract_sibling_refs(q, sibling_names): dependents[prereq].append(name) @@ -404,21 +404,21 @@ def _build_dependency_graph( @staticmethod def _kahn_sort( - in_degree: Dict[str, int], - dependents: Dict[str, List[str]], - ) -> List[str]: + in_degree: dict[str, int], + dependents: dict[str, list[str]], + ) -> list[str]: """Topologically sort by Kahn's algorithm. Cycle → ValueError. Mutates ``in_degree`` in place; callers shouldn't reuse it. The frontier is kept sorted for deterministic output order across runs. """ - frontier: List[str] = sorted(n for n, d in in_degree.items() if d == 0) - sorted_names: List[str] = [] + frontier: list[str] = sorted(n for n, d in in_degree.items() if d == 0) + sorted_names: list[str] = [] while frontier: n = frontier.pop(0) sorted_names.append(n) - unlocked: List[str] = [] + unlocked: list[str] = [] for dep in dependents[n]: in_degree[dep] -= 1 if in_degree[dep] == 0: @@ -435,8 +435,8 @@ def _kahn_sort( @classmethod def _topologically_order_queries( cls, - queries: List["SlayerQuery"], - ) -> List["SlayerQuery"]: + queries: list["SlayerQuery"], + ) -> list["SlayerQuery"]: """Re-order a runtime query list so every stage appears after the siblings it references via ``source_model`` or ``joins.target_model``. Lets callers submit a DAG in any order — @@ -477,11 +477,11 @@ def _topologically_order_queries( async def execute( # NOSONAR S3776 — public dispatch over str/dict/list/SlayerQuery; splitting hides the input-shape contract self, query: "SlayerQuery | dict | list[SlayerQuery | dict] | str", - variables: Optional[Dict[str, Any]] = None, + variables: dict[str, Any] | None = None, *, dry_run: bool = False, explain: bool = False, - data_source: Optional[str] = None, + data_source: str | None = None, ) -> SlayerResponse: runtime_kwarg = variables or {} @@ -536,10 +536,10 @@ async def execute( # NOSONAR S3776 — public dispatch over str/dict/list/Slaye async def _execute_by_name( self, name: str, - runtime_kwarg: Dict[str, Any], + runtime_kwarg: dict[str, Any], dry_run: bool = False, explain: bool = False, - data_source: Optional[str] = None, + data_source: str | None = None, ) -> SlayerResponse: """Run the backing query of a query-backed model by name.""" model = await self.storage.get_model(name, data_source=data_source) @@ -553,7 +553,7 @@ async def _execute_by_name( stages = list(model.source_queries) main_query = stages[-1] - named_queries: Dict[str, SlayerQuery] = {} + named_queries: dict[str, SlayerQuery] = {} for q in stages[:-1]: if q.name: if q.name in named_queries: @@ -587,12 +587,12 @@ async def _execute_by_name( async def _execute_pipeline( # NOSONAR S3776 — linear pipeline (resolve→enrich→generate→execute); breaking it up obscures the order of operations self, query: SlayerQuery, - named_queries: Dict[str, SlayerQuery], - runtime_kwarg: Dict[str, Any], + named_queries: dict[str, SlayerQuery], + runtime_kwarg: dict[str, Any], *, dry_run: bool = False, explain: bool = False, - prefer_data_source: Optional[str] = None, + prefer_data_source: str | None = None, ) -> SlayerResponse: """Shared pipeline used by both ``execute()`` and ``_execute_by_name()``. @@ -648,8 +648,8 @@ async def _execute_pipeline( # NOSONAR S3776 — linear pipeline (resolve→enr # entry is included only if its alias is part of the public # projection (filter-extracted hidden transforms, ORDER-BY # aggregates, and window-arg hoists are silently dropped). - dim_meta: Dict[str, FieldMetadata] = {} - measure_meta: Dict[str, FieldMetadata] = {} + dim_meta: dict[str, FieldMetadata] = {} + measure_meta: dict[str, FieldMetadata] = {} for d in enriched.dimensions: if d.alias in public_aliases and (d.label or d.format): dim_meta[d.alias] = FieldMetadata(label=d.label, format=d.format) @@ -805,7 +805,7 @@ async def _collect_models_touched( return touched async def _expand_join_graph( - self, *, touched: "set[str]", data_source: Optional[str] + self, *, touched: "set[str]", data_source: str | None ) -> None: """Follow each touched model's joins transitively, adding reachable target_model names to ``touched``. Visited-set guarded to avoid @@ -853,7 +853,7 @@ async def _maybe_raise_schema_drift( # time), so attribution only needs the parent's data_source. data_sources: set[str] = {model.data_source} if model.data_source else set() - collected: List[Any] = [] + collected: list[Any] = [] for ds_name in data_sources or {None}: try: entries = await self.validate_models(data_source=ds_name) @@ -891,7 +891,7 @@ def _build_type_probe_query(self, model: SlayerModel) -> SlayerQuery: types) and falls back to the first allowed aggregation otherwise. Skips primary-key columns (they're identifiers, not values to probe). """ - measures: List[ModelMeasure] = [] + measures: list[ModelMeasure] = [] for c in model.columns: if c.hidden or c.primary_key: continue @@ -908,8 +908,8 @@ def _build_type_probe_query(self, model: SlayerModel) -> SlayerQuery: async def get_column_types( self, model_name: str, - data_source: Optional[str] = None, - ) -> Dict[str, str]: + data_source: str | None = None, + ) -> dict[str, str]: """Infer column types for a model's columns via a type-probe query. Builds a real query through the engine's enrich+generate pipeline @@ -981,7 +981,7 @@ async def get_column_types( return {} # Map qualified aliases (e.g., "orders.revenue_max") back to bare measure names - result: Dict[str, str] = {} + result: dict[str, str] = {} for em in enriched.measures: if em.alias in raw_types: result[em.source_measure_name or em.name] = raw_types[em.alias] @@ -1003,7 +1003,7 @@ async def aclose(self) -> None: def execute_sync( self, query: "SlayerQuery | dict | list[SlayerQuery | dict] | str", - variables: Optional[Dict[str, Any]] = None, + variables: dict[str, Any] | None = None, *, dry_run: bool = False, explain: bool = False, @@ -1027,12 +1027,12 @@ async def edit_model_remove( self, *, model_name: str, - data_source: Optional[str], - remove_columns: Optional[List[str]] = None, - remove_measures: Optional[List[str]] = None, - remove_aggregations: Optional[List[str]] = None, - remove_joins: Optional[List[str]] = None, - remove_filters: Optional[List[str]] = None, + data_source: str | None, + remove_columns: list[str] | None = None, + remove_measures: list[str] | None = None, + remove_aggregations: list[str] | None = None, + remove_joins: list[str] | None = None, + remove_filters: list[str] | None = None, ) -> SlayerModel: """Apply surgical removals to a persisted model. @@ -1099,13 +1099,13 @@ async def edit_model_remove( return updated async def delete_model_by_name( - self, *, model_name: str, data_source: Optional[str] + self, *, model_name: str, data_source: str | None ) -> bool: """Delete a persisted model by name. Returns True if the model existed.""" return await self.storage.delete_model(model_name, data_source=data_source) async def apply_drift_deletes( - self, deletes: "List[Any]" + self, deletes: "list[Any]" ) -> "Any": """Apply each ``ToDeleteEntry`` via the engine helpers and return the combined ``ApplyDriftResult`` (applied, errors, residual). @@ -1122,8 +1122,8 @@ async def apply_drift_deletes( ApplyError, ) - applied: List[AppliedEntry] = [] - errors: List[ApplyError] = [] + applied: list[AppliedEntry] = [] + errors: list[ApplyError] = [] touched_ds: set[str] = set() for entry in deletes: @@ -1166,7 +1166,7 @@ async def apply_drift_deletes( ) # Re-validate the touched datasources to compute residual drift. - residual: List[Any] = [] + residual: list[Any] = [] for ds_name in touched_ds: try: residual.extend(await self.validate_models(data_source=ds_name)) @@ -1183,8 +1183,8 @@ async def apply_drift_deletes( ) async def validate_models( - self, data_source: Optional[str] = None - ) -> "List[Any]": + self, data_source: str | None = None + ) -> "list[Any]": """Diff persisted models against live database schemas. Returns the minimal list of deletes needed for SQL generation to @@ -1205,7 +1205,7 @@ async def validate_models( return [] identities = await self.storage._list_all_model_identities() ds_model_names = [n for d, n in identities if d == data_source] - models: List[SlayerModel] = [] + models: list[SlayerModel] = [] for name in ds_model_names: m = await self.storage.get_model(name, data_source=data_source) if m is not None: @@ -1220,13 +1220,13 @@ async def validate_models( if not ds_names: return [] - async def _validate_one(name: str) -> "List[ToDeleteEntry]": + async def _validate_one(name: str) -> "list[ToDeleteEntry]": return await self.validate_models(data_source=name) results = await _asyncio.gather( *(_validate_one(n) for n in ds_names), return_exceptions=True ) - out: List = [] + out: list = [] for r in results: if isinstance(r, BaseException): logger.warning("validate_models: per-DS validation failed: %s", r) @@ -1238,8 +1238,8 @@ def create_model_from_query_sync( self, query: "SlayerQuery | list[SlayerQuery] | dict | list[dict]", name: str, - description: Optional[str] = None, - variables: Optional[Dict[str, Any]] = None, + description: str | None = None, + variables: dict[str, Any] | None = None, save: bool = True, ) -> SlayerModel: """Synchronous wrapper for create_model_from_query().""" @@ -1258,10 +1258,10 @@ def create_model_from_query_sync( async def _expand_query_backed_model( self, model: SlayerModel, - outer_vars: Optional[Dict[str, Any]], - runtime_kwarg: Optional[Dict[str, Any]], + outer_vars: dict[str, Any] | None, + runtime_kwarg: dict[str, Any] | None, dry_run_placeholders: bool, - _resolving: Optional[set], + _resolving: set | None, ) -> SlayerModel: """If ``model`` is query-backed, expand its ``source_queries`` into a virtual model (with rendered SQL). Otherwise return ``model`` unchanged. @@ -1290,10 +1290,10 @@ async def _resolve_query_model( # NOSONAR S3776 — type-dispatch on str/Slayer query_model, named_queries: dict = None, _resolving: set = None, - outer_vars: Optional[Dict[str, Any]] = None, - runtime_kwarg: Optional[Dict[str, Any]] = None, + outer_vars: dict[str, Any] | None = None, + runtime_kwarg: dict[str, Any] | None = None, dry_run_placeholders: bool = False, - prefer_data_source: Optional[str] = None, + prefer_data_source: str | None = None, ) -> SlayerModel: """Resolve query.source_model — handles str, SlayerModel, and ModelExtension.""" from slayer.core.query import ModelExtension @@ -1377,10 +1377,10 @@ async def _resolve_model( model_name: str, named_queries: dict[str, SlayerQuery] = None, _resolving: set = None, - outer_vars: Optional[Dict[str, Any]] = None, - runtime_kwarg: Optional[Dict[str, Any]] = None, + outer_vars: dict[str, Any] | None = None, + runtime_kwarg: dict[str, Any] | None = None, dry_run_placeholders: bool = False, - prefer_data_source: Optional[str] = None, + prefer_data_source: str | None = None, ) -> SlayerModel: """Resolve a model by name — checks named queries first, then storage.""" named_queries = named_queries or {} @@ -1411,10 +1411,10 @@ async def _resolve_model_inner( model_name: str, named_queries: dict[str, SlayerQuery], _resolving: set = None, - outer_vars: Optional[Dict[str, Any]] = None, - runtime_kwarg: Optional[Dict[str, Any]] = None, + outer_vars: dict[str, Any] | None = None, + runtime_kwarg: dict[str, Any] | None = None, dry_run_placeholders: bool = False, - prefer_data_source: Optional[str] = None, + prefer_data_source: str | None = None, ) -> SlayerModel: # Named query overrides stored model if model_name in named_queries: @@ -1477,8 +1477,8 @@ async def create_model_from_query( self, query: "SlayerQuery | list[SlayerQuery] | dict | list[dict]", name: str, - description: Optional[str] = None, - variables: Optional[Dict[str, Any]] = None, + description: str | None = None, + variables: dict[str, Any] | None = None, save: bool = True, ) -> SlayerModel: """Create a query-backed model from a query (or list of stages). @@ -1527,7 +1527,7 @@ async def save_model(self, model: SlayerModel) -> SlayerModel: # up the old storage entry when a query-backed model's resolved # data_source changes (e.g. its backing query now points at a # different upstream datasource). - prior_data_source: Optional[str] = None + prior_data_source: str | None = None if model.source_queries: try: identity = await self.storage.resolve_model_identity(model.name) @@ -1596,7 +1596,7 @@ async def _enrich( # NOSONAR S3776 — orchestrates resolve-callback closures + query: SlayerQuery, model: SlayerModel, named_queries: dict[str, SlayerQuery] = None, - dialect: Optional[str] = None, + dialect: str | None = None, *, drop_unreachable_filters: bool = False, ) -> EnrichedQuery: @@ -1725,7 +1725,7 @@ async def _resolve_model_for_expansion(model_name, named_queries): async def _render_query_backed_join_target( self, target: SlayerModel, - outer_query_variables: Optional[Dict[str, Any]], + outer_query_variables: dict[str, Any] | None, ) -> SlayerModel: """Resolve a query-backed model used as a JOIN target. @@ -1774,8 +1774,8 @@ async def _query_as_model( # NOSONAR S3776 — variable-precedence + enrich + S named_queries: dict[str, SlayerQuery] = None, override_name: str = None, _resolving: set = None, - outer_vars: Optional[Dict[str, Any]] = None, - runtime_kwarg: Optional[Dict[str, Any]] = None, + outer_vars: dict[str, Any] | None = None, + runtime_kwarg: dict[str, Any] | None = None, dry_run_placeholders: bool = False, ) -> SlayerModel: """Build a virtual SlayerModel from a nested query's result. @@ -1814,7 +1814,7 @@ async def _query_as_model( # NOSONAR S3776 — variable-precedence + enrich + S scoped = self._scope_named_queries_to_prior( named_queries, inner_query.name ) - forbidden_now: Dict[str, str] = {} + forbidden_now: dict[str, str] = {} if scoped is not named_queries and inner_query.name: for k in named_queries: if k not in scoped: @@ -1969,7 +1969,7 @@ def _alias_to_short(alias: str) -> str: # One Column per result column — each is potentially both a dimension # (group-by) or measure (with colon-aggregation) at query time. - cols: List[Column] = [] + cols: list[Column] = [] for _, short, dtype, label, desc, fmt in column_map: cols.append(Column(name=short, sql=short, type=dtype, label=label, description=desc, format=fmt)) @@ -2128,7 +2128,7 @@ async def _auto_move_fields_to_dimensions( if not query.measures: return query - kept: List = [] + kept: list = [] extra_dims = list(query.dimensions or []) moved = False @@ -2517,7 +2517,7 @@ async def _resolve_datasource(self, model: SlayerModel) -> DatasourceConfig: return ds @staticmethod - def _dialect_for_type(ds_type: Optional[str]) -> str: + def _dialect_for_type(ds_type: str | None) -> str: """Map a datasource-config ``type`` string to a sqlglot dialect name. Delegates to ``dialect_for_ds_type`` in ``slayer.sql.dialects`` diff --git a/slayer/engine/schema_drift.py b/slayer/engine/schema_drift.py index fd0d0dcb..29e9e7b4 100644 --- a/slayer/engine/schema_drift.py +++ b/slayer/engine/schema_drift.py @@ -18,13 +18,7 @@ from typing import ( Annotated, Any, - Dict, - List, Literal, - Optional, - Set, - Tuple, - Union, ) import sqlalchemy as sa @@ -74,10 +68,10 @@ class DeleteReason(BaseModel): class RemoveSpec(BaseModel): """Per-entity removal spec, mirroring the MCP ``edit_model`` ``remove=`` shape.""" - columns: List[str] = Field(default_factory=list) - measures: List[str] = Field(default_factory=list) - aggregations: List[str] = Field(default_factory=list) - joins: List[str] = Field(default_factory=list) + columns: list[str] = Field(default_factory=list) + measures: list[str] = Field(default_factory=list) + aggregations: list[str] = Field(default_factory=list) + joins: list[str] = Field(default_factory=list) class EditModelDelete(BaseModel): @@ -87,8 +81,8 @@ class EditModelDelete(BaseModel): model_name: str data_source: str remove: RemoveSpec = Field(default_factory=RemoveSpec) - remove_filters: List[str] = Field(default_factory=list) - reasons: List[DeleteReason] = Field(default_factory=list) + remove_filters: list[str] = Field(default_factory=list) + reasons: list[DeleteReason] = Field(default_factory=list) class WholeModelDelete(BaseModel): @@ -97,11 +91,11 @@ class WholeModelDelete(BaseModel): tool: Literal["delete_model"] = "delete_model" model_name: str data_source: str - reasons: List[DeleteReason] = Field(default_factory=list) + reasons: list[DeleteReason] = Field(default_factory=list) ToDeleteEntry = Annotated[ - Union[EditModelDelete, WholeModelDelete], Field(discriminator="tool") + EditModelDelete | WholeModelDelete, Field(discriminator="tool") ] @@ -111,11 +105,11 @@ class ModelAddition(BaseModel): model_name: str data_source: str created: bool = False # True if the model was new - new_columns: List[str] = Field(default_factory=list) - new_joins: List[str] = Field(default_factory=list) + new_columns: list[str] = Field(default_factory=list) + new_joins: list[str] = Field(default_factory=list) # DEV-1538: persisted INT columns whose type widened (to DOUBLE or TEXT) # because the SQLite affinity probe disagreed with the declared type. - widened_columns: List[str] = Field(default_factory=list) + widened_columns: list[str] = Field(default_factory=list) class IngestionError(BaseModel): @@ -129,9 +123,9 @@ class IngestionError(BaseModel): class IdempotentIngestResult(BaseModel): """Combined return shape of the idempotent ``ingest_datasource`` pass.""" - additions: List[ModelAddition] = Field(default_factory=list) - to_delete: List[ToDeleteEntry] = Field(default_factory=list) - errors: List[IngestionError] = Field(default_factory=list) + additions: list[ModelAddition] = Field(default_factory=list) + to_delete: list[ToDeleteEntry] = Field(default_factory=list) + errors: list[IngestionError] = Field(default_factory=list) class AppliedEntry(BaseModel): @@ -154,9 +148,9 @@ class ApplyError(BaseModel): class ApplyDriftResult(BaseModel): """Combined return shape of ``apply_drift_deletes``.""" - applied: List[AppliedEntry] = Field(default_factory=list) - errors: List[ApplyError] = Field(default_factory=list) - residual: List[ToDeleteEntry] = Field(default_factory=list) + applied: list[AppliedEntry] = Field(default_factory=list) + errors: list[ApplyError] = Field(default_factory=list) + residual: list[ToDeleteEntry] = Field(default_factory=list) # =========================================================================== @@ -173,10 +167,10 @@ class LiveTable(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - columns: Dict[str, DataType] = Field(default_factory=dict) - pk_columns: Set[str] = Field(default_factory=set) + columns: dict[str, DataType] = Field(default_factory=dict) + pk_columns: set[str] = Field(default_factory=set) # Each entry: (local_column, ref_table, ref_column) - fk_relationships: List[Tuple[str, str, str]] = Field(default_factory=list) + fk_relationships: list[tuple[str, str, str]] = Field(default_factory=list) # =========================================================================== @@ -205,7 +199,7 @@ def data_type_bucket(dt: DataType) -> str: return str(dt) -def _is_bare_identifier(s: Optional[str]) -> bool: +def _is_bare_identifier(s: str | None) -> bool: """``s`` is a bare SQL identifier (alphanumeric + underscore, no leading digit).""" if not s: return False @@ -215,7 +209,7 @@ def _is_bare_identifier(s: Optional[str]) -> bool: return all(c.isalnum() or c == "_" for c in s) -def _column_is_base(col_sql: Optional[str]) -> bool: +def _column_is_base(col_sql: str | None) -> bool: """A Column whose ``sql`` is None or a bare identifier is a "base" column — it claims a live column. Derived expressions (``amount * 2``, ``customers.region``, etc.) do not. @@ -232,10 +226,10 @@ def _column_is_base(col_sql: Optional[str]) -> bool: def _diff_sql_table_columns( *, model: SlayerModel, live_table: LiveTable -) -> Tuple[List[str], List[DeleteReason]]: +) -> tuple[list[str], list[DeleteReason]]: """Per-column diff of a sql_table-mode model against live columns.""" - dropped: List[str] = [] - reasons: List[DeleteReason] = [] + dropped: list[str] = [] + reasons: list[DeleteReason] = [] for col in model.columns: # Only compare base columns directly. Derived columns are handled # by cascade. @@ -270,12 +264,12 @@ def _diff_sql_table_joins( *, model: SlayerModel, live_table: LiveTable, - available_models_in_ds: Set[str], -) -> Tuple[List[str], List[DeleteReason]]: + available_models_in_ds: set[str], +) -> tuple[list[str], list[DeleteReason]]: """Per-join diff of a sql_table-mode model against live FK columns and in-datasource model availability.""" - dropped: List[str] = [] - reasons: List[DeleteReason] = [] + dropped: list[str] = [] + reasons: list[DeleteReason] = [] # ``join.join_pairs[*][0]`` is the semantic column name (``Column.name``). # Resolve to the physical column name via ``Column.sql`` before checking # against the live table — for a base column like @@ -322,9 +316,9 @@ def _diff_sql_table_joins( def diff_sql_table_model( *, model: SlayerModel, - live_table: Optional[LiveTable], - available_models_in_ds: Set[str], -) -> Tuple[Optional[ToDeleteEntry], Set[str]]: + live_table: LiveTable | None, + available_models_in_ds: set[str], +) -> tuple[ToDeleteEntry | None, set[str]]: """Diff a sql_table-mode model against live introspection. Returns ``(entry_or_None, dropped_column_names)``. @@ -384,8 +378,8 @@ def diff_sql_table_model( def diff_sql_model( *, model: SlayerModel, - live_columns: Optional[Dict[str, DataType]], -) -> Tuple[Optional[ToDeleteEntry], Set[str]]: + live_columns: dict[str, DataType] | None, +) -> tuple[ToDeleteEntry | None, set[str]]: """Diff a sql-mode model against trial-execute cursor metadata. ``live_columns is None`` ⇒ trial-execute failed ⇒ ``WholeModelDelete``. @@ -409,8 +403,8 @@ def diff_sql_model( {c.name for c in model.columns}, ) - dropped_cols: List[str] = [] - reasons: List[DeleteReason] = [] + dropped_cols: list[str] = [] + reasons: list[DeleteReason] = [] for col in model.columns: # Cursor exposes ALIAS names — match by col.name first, fall back to # col.sql for legacy cases where a Column's name differs from its @@ -463,7 +457,7 @@ def diff_sql_model( # =========================================================================== -def _extract_column_refs_from_sql(sql: str) -> List[Tuple[Optional[str], str]]: +def _extract_column_refs_from_sql(sql: str) -> list[tuple[str | None, str]]: """Return all ``(table_alias, column_name)`` refs in a SQL expression. ``table_alias`` is ``None`` for bare identifiers, the raw alias string @@ -475,7 +469,7 @@ def _extract_column_refs_from_sql(sql: str) -> List[Tuple[Optional[str], str]]: parsed = sqlglot.parse_one(sql) except Exception: return [] - refs: List[Tuple[Optional[str], str]] = [] + refs: list[tuple[str | None, str]] = [] for col in parsed.find_all(exp.Column): if col.args.get("db") or col.args.get("catalog"): continue @@ -485,21 +479,21 @@ def _extract_column_refs_from_sql(sql: str) -> List[Tuple[Optional[str], str]]: return refs -def _agg_ref_names(agg_refs: Dict[str, AggregatedMeasureRef]) -> Set[str]: +def _agg_ref_names(agg_refs: dict[str, AggregatedMeasureRef]) -> set[str]: """Names from a ``measure:agg`` placeholder map, excluding ``*``.""" return {ref.measure_name for ref in agg_refs.values() if ref.measure_name != "*"} def _bare_measure_names( - measure_names: List[str], - agg_refs: Dict[str, AggregatedMeasureRef], + measure_names: list[str], + agg_refs: dict[str, AggregatedMeasureRef], *, - skip_placeholder_prefix: Optional[str] = None, -) -> Set[str]: + skip_placeholder_prefix: str | None = None, +) -> set[str]: """Filter raw ``measure_names`` to the ones that are not colon-syntax placeholders, optionally also stripping sub-transform placeholders. """ - out: Set[str] = set() + out: set[str] = set() for n in measure_names: if n in agg_refs: continue @@ -509,7 +503,7 @@ def _bare_measure_names( return out -def _walk_field_spec_measure_refs(spec: Any) -> Set[str]: +def _walk_field_spec_measure_refs(spec: Any) -> set[str]: """Walk a ``FieldSpec`` (parse_formula output) and return the set of measure_name strings (which may be dotted: ``"customers.revenue"``). """ @@ -534,8 +528,8 @@ def _walk_field_spec_measure_refs(spec: Any) -> Set[str]: def _measure_formula_refs( formula: str, *, - named_measures: Optional[Dict[str, str]] = None, -) -> Set[str]: + named_measures: dict[str, str] | None = None, +) -> set[str]: """Best-effort: parse ``formula`` and return the set of column / measure names it references. Returns the empty set on any parse failure. @@ -551,7 +545,7 @@ def _measure_formula_refs( return _walk_field_spec_measure_refs(spec) -def _filter_refs(filter_str: str) -> List[str]: +def _filter_refs(filter_str: str) -> list[str]: """Best-effort: return list of column references in a SQL-mode filter. Used to scan ``Column.filter`` / ``SlayerModel.filters`` strings (Mode A @@ -564,7 +558,7 @@ def _filter_refs(filter_str: str) -> List[str]: return list(pf.columns) -def _filter_refs_dsl(filter_str: str) -> List[str]: +def _filter_refs_dsl(filter_str: str) -> list[str]: """Best-effort: return list of column / measure references in a DSL filter. Used to scan ``SlayerQuery.filters`` strings (Mode B DSL — DEV-1369), @@ -592,8 +586,8 @@ def _walk_alias_to_target_model( *, source_model: SlayerModel, table_alias: str, - models_by_name: Dict[str, SlayerModel], -) -> Optional[SlayerModel]: + models_by_name: dict[str, SlayerModel], +) -> SlayerModel | None: """Resolve a ``__``-delimited path alias starting from ``source_model`` to the terminal joined model. Returns None if any hop fails. """ @@ -616,8 +610,8 @@ def _resolve_dotted_ref_to_model( *, source_model: SlayerModel, dotted_ref: str, - models_by_name: Dict[str, SlayerModel], -) -> Tuple[Optional[SlayerModel], Optional[str]]: + models_by_name: dict[str, SlayerModel], +) -> tuple[SlayerModel | None, str | None]: """Resolve a dotted measure/column ref like ``customers.region`` or ``customers.regions.name`` to ``(target_model, leaf_name)``. @@ -639,13 +633,13 @@ def _resolve_dotted_ref_to_model( # =========================================================================== -def _pk_columns(model: SlayerModel) -> Set[str]: +def _pk_columns(model: SlayerModel) -> set[str]: return {c.name for c in model.columns if c.primary_key} def _ensure_edit_entry( *, - edit_entries: Dict[str, EditModelDelete], + edit_entries: dict[str, EditModelDelete], model: SlayerModel, ) -> EditModelDelete: if model.name not in edit_entries: @@ -658,8 +652,8 @@ def _ensure_edit_entry( def _add_dropped_column( *, - edit_entries: Dict[str, EditModelDelete], - dropped_cols: Dict[str, Set[str]], + edit_entries: dict[str, EditModelDelete], + dropped_cols: dict[str, set[str]], model: SlayerModel, column_name: str, reason: str, @@ -679,8 +673,8 @@ def _add_dropped_column( def _add_dropped_measure( *, - edit_entries: Dict[str, EditModelDelete], - dropped_measures: Dict[str, Set[str]], + edit_entries: dict[str, EditModelDelete], + dropped_measures: dict[str, set[str]], model: SlayerModel, measure_name: str, reason: str, @@ -699,8 +693,8 @@ def _add_dropped_measure( def _add_dropped_join( *, - edit_entries: Dict[str, EditModelDelete], - dropped_joins: Dict[str, Set[str]], + edit_entries: dict[str, EditModelDelete], + dropped_joins: dict[str, set[str]], model: SlayerModel, target_name: str, reason: str, @@ -719,7 +713,7 @@ def _add_dropped_join( def _add_remove_filter( *, - edit_entries: Dict[str, EditModelDelete], + edit_entries: dict[str, EditModelDelete], model: SlayerModel, filter_text: str, reason: str, @@ -742,8 +736,8 @@ def _add_remove_filter( def _resolve_stage_source_to_base( *, source_model: object, - prior_stages_by_name: Dict[str, SlayerQuery], -) -> Optional[str]: + prior_stages_by_name: dict[str, SlayerQuery], +) -> str | None: """Walk a ``source_model`` reference (str / SlayerModel / ModelExtension / prior-stage-name) back to a real persisted base model name. @@ -751,7 +745,7 @@ def _resolve_stage_source_to_base( underlying model — we follow it transparently so query-backed drift attribution doesn't silently skip extension-wrapped stages. """ - seen: Set[str] = set() + seen: set[str] = set() current = source_model while True: if isinstance(current, str): @@ -784,29 +778,29 @@ class _StageGraph(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - stage_source_name: Optional[str] = None - extension_targets: Set[str] = Field(default_factory=set) - reachable: Set[str] = Field(default_factory=set) - models_by_name: Dict[str, SlayerModel] = Field(default_factory=dict) + stage_source_name: str | None = None + extension_targets: set[str] = Field(default_factory=set) + reachable: set[str] = Field(default_factory=set) + models_by_name: dict[str, SlayerModel] = Field(default_factory=dict) def _build_stage_graph( *, stage: SlayerQuery, - stage_source_name: Optional[str], - models_by_name: Dict[str, SlayerModel], + stage_source_name: str | None, + models_by_name: dict[str, SlayerModel], ) -> _StageGraph: """Build a ``_StageGraph`` for a single stage. ``stage_source_name`` is the resolved base model name (str), or ``None`` for inline / unresolved sources. """ extension_targets = _stage_join_targets(stage) - reachable: Set[str] = set() + reachable: set[str] = set() if stage_source_name: reachable.add(stage_source_name) reachable |= extension_targets frontier = list(reachable) - visited: Set[str] = set() + visited: set[str] = set() while frontier: name = frontier.pop() if name in visited: @@ -832,7 +826,7 @@ def _attribute_ref_to_base( ref: str, base_name: str, graph: _StageGraph, -) -> Optional[str]: +) -> str | None: """Walk ``ref`` through the stage's join graph and return the leaf column name when it resolves to ``base_name``, else ``None``. @@ -868,8 +862,8 @@ def _attribute_ref_to_base( def _measure_refs_on_base( stage: SlayerQuery, base_name: str, graph: _StageGraph -) -> Set[str]: - out: Set[str] = set() +) -> set[str]: + out: set[str] = set() for m in stage.measures or []: formula = getattr(m, "formula", None) if not formula: @@ -885,8 +879,8 @@ def _measure_refs_on_base( def _dimension_refs_on_base( stage: SlayerQuery, base_name: str, graph: _StageGraph -) -> Set[str]: - out: Set[str] = set() +) -> set[str]: + out: set[str] = set() for d in stage.dimensions or []: full = getattr(d, "full_name", None) or str(d) attributed = _attribute_ref_to_base( @@ -899,8 +893,8 @@ def _dimension_refs_on_base( def _time_dimension_refs_on_base( stage: SlayerQuery, base_name: str, graph: _StageGraph -) -> Set[str]: - out: Set[str] = set() +) -> set[str]: + out: set[str] = set() for td in stage.time_dimensions or []: attributed = _attribute_ref_to_base( ref=td.dimension.full_name, base_name=base_name, graph=graph @@ -912,8 +906,8 @@ def _time_dimension_refs_on_base( def _filter_refs_on_base( stage: SlayerQuery, base_name: str, graph: _StageGraph -) -> Set[str]: - out: Set[str] = set() +) -> set[str]: + out: set[str] = set() # ``SlayerQuery.filters`` are Mode B (DSL) — go through the DSL parser # so colon-syntax aggregations and transforms surface their underlying # measure names. ``_filter_refs`` (SQL-mode) would drop them silently. @@ -931,8 +925,8 @@ def _stage_referenced_columns_for_base( *, stage: SlayerQuery, base_name: str, - graph: Optional[_StageGraph] = None, -) -> Set[str]: + graph: _StageGraph | None = None, +) -> set[str]: """Return the set of column names referenced *on* ``base_name`` by a single source_queries stage. Walks the stage's join graph (passed via ``graph``) so multi-hop dotted refs and ModelExtension-added joins are @@ -957,7 +951,7 @@ def _stage_referenced_columns_for_base( ) -def _stage_join_targets(stage: SlayerQuery) -> Set[str]: +def _stage_join_targets(stage: SlayerQuery) -> set[str]: """Return the set of join target_model names referenced by a stage. ``SlayerQuery`` itself has no ``joins`` field; joins on a stage live @@ -968,7 +962,7 @@ def _stage_join_targets(stage: SlayerQuery) -> Set[str]: """ source = getattr(stage, "source_model", None) joins = getattr(source, "joins", None) or [] - out: Set[str] = set() + out: set[str] = set() for j in joins: target = getattr(j, "target_model", None) if isinstance(target, str): @@ -981,9 +975,9 @@ def _check_stage_against_base( stage: SlayerQuery, base_name: str, graph: _StageGraph, - dropped_cols: Dict[str, Set[str]], - pk_per_model: Dict[str, Set[str]], -) -> Set[str]: + dropped_cols: dict[str, set[str]], + pk_per_model: dict[str, set[str]], +) -> set[str]: """Return the set of dropped column names on ``base_name`` that this stage references (resolved through the stage's join graph). @@ -1004,8 +998,8 @@ def _stage_uses_dropped_join( stage: SlayerQuery, base_name: str, graph: _StageGraph, - dropped_joins: Dict[str, Set[str]], -) -> Optional[str]: + dropped_joins: dict[str, set[str]], +) -> str | None: """If ``stage`` references any column under a join target that's been dropped on ``base_name``, return the conflicting target; else None. @@ -1036,12 +1030,12 @@ def _check_stage_for_whole_drop( base_name: str, qb_name: str, graph: _StageGraph, - whole_dropped_models: Set[str], - dropped_cols: Dict[str, Set[str]], - dropped_joins: Dict[str, Set[str]], - pk_per_model: Dict[str, Set[str]], - candidate_base_names: Set[str], -) -> Optional[DeleteReason]: + whole_dropped_models: set[str], + dropped_cols: dict[str, set[str]], + dropped_joins: dict[str, set[str]], + pk_per_model: dict[str, set[str]], + candidate_base_names: set[str], +) -> DeleteReason | None: """Decide whether a single stage of a query-backed model triggers the whole-drop. Returns a ``DeleteReason`` on the first matching trigger, or ``None`` when the stage has no fatal references. @@ -1099,13 +1093,13 @@ def _check_stage_for_whole_drop( def _query_backed_should_whole_drop( *, qb_model: SlayerModel, - dropped_cols: Dict[str, Set[str]], - dropped_joins: Dict[str, Set[str]], - whole_dropped_models: Set[str], - pk_per_model: Dict[str, Set[str]], - candidate_base_names: Optional[Set[str]] = None, - models_by_name: Optional[Dict[str, SlayerModel]] = None, -) -> Optional[DeleteReason]: + dropped_cols: dict[str, set[str]], + dropped_joins: dict[str, set[str]], + whole_dropped_models: set[str], + pk_per_model: dict[str, set[str]], + candidate_base_names: set[str] | None = None, + models_by_name: dict[str, SlayerModel] | None = None, +) -> DeleteReason | None: """Return a non-None DeleteReason when this query-backed model should be whole-dropped due to cascading from base-model drift, else None. @@ -1122,7 +1116,7 @@ def _query_backed_should_whole_drop( models_by_name = models_by_name or {} for i, stage in enumerate(stages): - prior_by_name: Dict[str, SlayerQuery] = {} + prior_by_name: dict[str, SlayerQuery] = {} for s in stages[:i]: s_name = getattr(s, "name", None) if s_name: @@ -1180,13 +1174,13 @@ class _CascadeState: def __init__( self, *, - models_by_name: Dict[str, SlayerModel], - edit_entries: Dict[str, EditModelDelete], - whole_entries: Dict[str, WholeModelDelete], - dropped_cols: Dict[str, Set[str]], - dropped_measures: Dict[str, Set[str]], - dropped_joins: Dict[str, Set[str]], - pk_per_model: Dict[str, Set[str]], + models_by_name: dict[str, SlayerModel], + edit_entries: dict[str, EditModelDelete], + whole_entries: dict[str, WholeModelDelete], + dropped_cols: dict[str, set[str]], + dropped_measures: dict[str, set[str]], + dropped_joins: dict[str, set[str]], + pk_per_model: dict[str, set[str]], ) -> None: self.models_by_name = models_by_name self.edit_entries = edit_entries @@ -1196,18 +1190,18 @@ def __init__( self.dropped_joins = dropped_joins self.pk_per_model = pk_per_model - def cascadable(self, name: str) -> Set[str]: + def cascadable(self, name: str) -> set[str]: """Cascadable column drops on ``name`` (excludes PKs — rule 7).""" return self.dropped_cols.get(name, set()) - self.pk_per_model.get(name, set()) def _column_ref_targets_dropped( *, - table_alias: Optional[str], + table_alias: str | None, ref_col: str, model: SlayerModel, state: _CascadeState, -) -> Tuple[bool, Optional[SlayerModel]]: +) -> tuple[bool, SlayerModel | None]: """Decide if a single ``(table_alias, ref_col)`` reference resolves to a dropped column. Returns ``(is_dropped, resolved_target_model)``. """ @@ -1225,7 +1219,7 @@ def _column_ref_targets_dropped( def _first_dropped_sql_column_ref( *, col: Column, model: SlayerModel, state: _CascadeState -) -> Optional[Tuple[SlayerModel, str]]: +) -> tuple[SlayerModel, str] | None: """Return ``(target_model, ref_col)`` for the first reference in ``col.sql`` that resolves to a dropped column, or ``None`` when nothing in the column's SQL references a dropped target. @@ -1272,7 +1266,7 @@ def _cascade_derived_columns( def _measure_drop_cause( *, ref: str, model: SlayerModel, state: _CascadeState -) -> Optional[str]: +) -> str | None: """If the measure ref resolves to a dropped column or measure, return a reason string; otherwise None. """ @@ -1292,10 +1286,10 @@ def _measure_drop_cause( def _first_dropped_cause( *, - refs: Set[str], + refs: set[str], model: SlayerModel, state: _CascadeState, -) -> Optional[str]: +) -> str | None: """Return the cause string for the first ref that resolves to a dropped column or measure, or ``None`` when nothing in ``refs`` is dropped. """ @@ -1409,7 +1403,7 @@ def _cascade_filters(*, model: SlayerModel, state: _CascadeState) -> bool: def _cascade_query_backed( - *, models: List[SlayerModel], state: _CascadeState + *, models: list[SlayerModel], state: _CascadeState ) -> bool: """Rule 6: query-backed model whose source_queries chain transitively references dropped state — whole-drop.""" @@ -1446,14 +1440,14 @@ def _cascade_query_backed( def _cascade_one_pass( *, - models: List[SlayerModel], - models_by_name: Dict[str, SlayerModel], - edit_entries: Dict[str, EditModelDelete], - whole_entries: Dict[str, WholeModelDelete], - dropped_cols: Dict[str, Set[str]], - dropped_measures: Dict[str, Set[str]], - dropped_joins: Dict[str, Set[str]], - pk_per_model: Dict[str, Set[str]], + models: list[SlayerModel], + models_by_name: dict[str, SlayerModel], + edit_entries: dict[str, EditModelDelete], + whole_entries: dict[str, WholeModelDelete], + dropped_cols: dict[str, set[str]], + dropped_measures: dict[str, set[str]], + dropped_joins: dict[str, set[str]], + pk_per_model: dict[str, set[str]], ) -> bool: """Run a single cascade pass; return True if anything new was added. @@ -1492,13 +1486,13 @@ def _cascade_one_pass( def _seed_one_diff_entry( *, model_name: str, - entry: Optional[ToDeleteEntry], - cols: Set[str], - edit_entries: Dict[str, EditModelDelete], - whole_entries: Dict[str, WholeModelDelete], - dropped_cols: Dict[str, Set[str]], - dropped_measures: Dict[str, Set[str]], - dropped_joins: Dict[str, Set[str]], + entry: ToDeleteEntry | None, + cols: set[str], + edit_entries: dict[str, EditModelDelete], + whole_entries: dict[str, WholeModelDelete], + dropped_cols: dict[str, set[str]], + dropped_measures: dict[str, set[str]], + dropped_joins: dict[str, set[str]], ) -> None: """Apply one ``(entry, dropped_columns)`` diff result to the cascade state dicts.""" @@ -1520,14 +1514,14 @@ def _seed_one_diff_entry( def _seed_state_from_diffs( *, - diffs_iterables: Tuple[ - Dict[str, Tuple[Optional[ToDeleteEntry], Set[str]]], ... + diffs_iterables: tuple[ + dict[str, tuple[ToDeleteEntry | None, set[str]]], ... ], - edit_entries: Dict[str, EditModelDelete], - whole_entries: Dict[str, WholeModelDelete], - dropped_cols: Dict[str, Set[str]], - dropped_measures: Dict[str, Set[str]], - dropped_joins: Dict[str, Set[str]], + edit_entries: dict[str, EditModelDelete], + whole_entries: dict[str, WholeModelDelete], + dropped_cols: dict[str, set[str]], + dropped_measures: dict[str, set[str]], + dropped_joins: dict[str, set[str]], ) -> None: """Populate the cascade state dicts from the base per-model diffs.""" for diffs in diffs_iterables: @@ -1546,13 +1540,13 @@ def _seed_state_from_diffs( def _collapse_entries( *, - edit_entries: Dict[str, EditModelDelete], - whole_entries: Dict[str, WholeModelDelete], -) -> List[ToDeleteEntry]: + edit_entries: dict[str, EditModelDelete], + whole_entries: dict[str, WholeModelDelete], +) -> list[ToDeleteEntry]: """Apply the collapse rule (whole-drop preempts edit on the same model) and return the final, name-sorted list of delete entries. """ - final: List[ToDeleteEntry] = [] + final: list[ToDeleteEntry] = [] for name in sorted(set(edit_entries.keys()) | set(whole_entries.keys())): if name in whole_entries: final.append(whole_entries[name]) @@ -1563,21 +1557,21 @@ def _collapse_entries( def compute_datasource_drops( *, - models: List[SlayerModel], - sql_table_diffs: Dict[str, Tuple[Optional[ToDeleteEntry], Set[str]]], - sql_diffs: Dict[str, Tuple[Optional[ToDeleteEntry], Set[str]]], -) -> List[ToDeleteEntry]: + models: list[SlayerModel], + sql_table_diffs: dict[str, tuple[ToDeleteEntry | None, set[str]]], + sql_diffs: dict[str, tuple[ToDeleteEntry | None, set[str]]], +) -> list[ToDeleteEntry]: """Combine per-model base diffs with cascade walking and the collapse rule. Pure: takes pre-computed diffs as input and returns the final flat list. Caller is responsible for restricting ``models`` to a single datasource — cascade walking does not cross datasource boundaries. """ - edit_entries: Dict[str, EditModelDelete] = {} - whole_entries: Dict[str, WholeModelDelete] = {} - dropped_cols: Dict[str, Set[str]] = {} - dropped_measures: Dict[str, Set[str]] = {} - dropped_joins: Dict[str, Set[str]] = {} + edit_entries: dict[str, EditModelDelete] = {} + whole_entries: dict[str, WholeModelDelete] = {} + dropped_cols: dict[str, set[str]] = {} + dropped_measures: dict[str, set[str]] = {} + dropped_joins: dict[str, set[str]] = {} _seed_state_from_diffs( diffs_iterables=(sql_table_diffs, sql_diffs), @@ -1619,8 +1613,8 @@ def compute_datasource_drops( def _live_schema_for_datasource( *, datasource: DatasourceConfig, - schema: Optional[str] = None, -) -> Dict[str, LiveTable]: + schema: str | None = None, +) -> dict[str, LiveTable]: """Return ``{table_name: LiveTable}`` for every live table in the DS, using SQLAlchemy ``Inspector`` and the same fallback path as auto-ingestion (``slayer/engine/ingestion.py``). @@ -1630,7 +1624,7 @@ def _live_schema_for_datasource( try: inspector = sa.inspect(sa_engine) table_names = list(inspector.get_table_names(schema=schema)) - out: Dict[str, LiveTable] = {} + out: dict[str, LiveTable] = {} for table_name in table_names: try: out[table_name] = _introspect_one_table( @@ -1661,7 +1655,7 @@ def _introspect_one_table( inspector: sa.engine.Inspector, sa_engine: sa.Engine, table_name: str, - schema: Optional[str], + schema: str | None, ) -> LiveTable: """Build a ``LiveTable`` for one table via the existing safe-introspection path used by ``slayer/engine/ingestion.py``. @@ -1670,7 +1664,7 @@ def _introspect_one_table( pk = _safe_get_pk_constraint(inspector, sa_engine, table_name, schema) pk_columns = set(pk.get("constrained_columns", []) or []) - columns: Dict[str, DataType] = {} + columns: dict[str, DataType] = {} for col in cols_meta: col_type = col["type"] if isinstance(col_type, DataType): @@ -1681,7 +1675,7 @@ def _introspect_one_table( # collapse to NUMBER. _ = _sa_type_is_float(col_type) - fks: List[Tuple[str, str, str]] = [] + fks: list[tuple[str, str, str]] = [] try: for fk in inspector.get_foreign_keys(table_name, schema=schema): constrained = fk.get("constrained_columns") or [] @@ -1714,7 +1708,7 @@ async def _live_columns_for_sql_model( *, model: SlayerModel, client: SlayerSQLClient, -) -> Optional[Dict[str, DataType]]: +) -> dict[str, DataType] | None: """Trial-execute ``model.sql`` with a 0-row guard and return cursor types. Returns ``None`` when the trial-execute itself fails — callers map that @@ -1762,8 +1756,8 @@ def _strip_ident_quotes(ident: str) -> str: def _resolve_live_table( - *, sql_table: str, live_tables: Dict[str, LiveTable] -) -> Optional[LiveTable]: + *, sql_table: str, live_tables: dict[str, LiveTable] +) -> LiveTable | None: """Look up a model's ``sql_table`` in the live introspection map, falling back to the bare name when the persisted value is schema- qualified (``schema.table``) and unquoting double-quoted identifiers @@ -1799,8 +1793,8 @@ def _is_validate_models_base_column(col: Column) -> bool: def _probe_validate_models_column( *, conn, model: SlayerModel, col: Column, table_name: str, - schema_name: Optional[str], -) -> Optional[DataType]: + schema_name: str | None, +) -> DataType | None: """Run the affinity probe for one column in a validate_models pass.""" from slayer.sql.sqlite_introspect import probe_sqlite_integer_column try: @@ -1835,8 +1829,8 @@ def _sqlite_probe_int_drift_for_model( *, model: SlayerModel, sa_engine, - default_schema: Optional[str] = None, -) -> List[Tuple[str, DataType, DeleteReason]]: + default_schema: str | None = None, +) -> list[tuple[str, DataType, DeleteReason]]: """DEV-1538: probe-driven type-drift detection on SQLite. For every persisted base column with ``Column.type == DataType.INT``, @@ -1864,7 +1858,7 @@ def _sqlite_probe_int_drift_for_model( else: schema_name, table_name = default_schema or None, model.sql_table - drifts: List[Tuple[str, DataType, DeleteReason]] = [] + drifts: list[tuple[str, DataType, DeleteReason]] = [] with sa_engine.connect() as conn: for col in model.columns: if not _is_validate_models_base_column(col): @@ -1886,8 +1880,8 @@ def _sqlite_probe_int_drift_for_model( async def _sqlite_probe_drifts_for_models( *, datasource: DatasourceConfig, - sql_table_models: List[SlayerModel], -) -> Dict[str, List[Tuple[str, DataType, DeleteReason]]]: + sql_table_models: list[SlayerModel], +) -> dict[str, list[tuple[str, DataType, DeleteReason]]]: """Run the SQLite probe for every model in one synchronous worker so the engine + connection lifecycle is shared across the validate pass.""" if not sql_table_models: @@ -1895,11 +1889,11 @@ async def _sqlite_probe_drifts_for_models( if (datasource.type or "").lower() != "sqlite": return {m.name: [] for m in sql_table_models} - def _run() -> Dict[str, List[Tuple[str, DataType, DeleteReason]]]: + def _run() -> dict[str, list[tuple[str, DataType, DeleteReason]]]: from slayer.sql import engine_factory sa_engine = engine_factory.get_engine(datasource.resolve_env_vars()) try: - out: Dict[str, List[Tuple[str, DataType, DeleteReason]]] = {} + out: dict[str, list[tuple[str, DataType, DeleteReason]]] = {} default_schema = datasource.schema_name or None for m in sql_table_models: out[m.name] = _sqlite_probe_int_drift_for_model( @@ -1918,9 +1912,9 @@ def _run() -> Dict[str, List[Tuple[str, DataType, DeleteReason]]]: def _merge_probe_drifts_into_diff( *, model: SlayerModel, - base_diff: Tuple[Optional[ToDeleteEntry], Set[str]], - probe_drifts: List[Tuple[str, DataType, DeleteReason]], -) -> Tuple[Optional[ToDeleteEntry], Set[str]]: + base_diff: tuple[ToDeleteEntry | None, set[str]], + probe_drifts: list[tuple[str, DataType, DeleteReason]], +) -> tuple[ToDeleteEntry | None, set[str]]: """Merge ``probe_drifts`` from :func:`_sqlite_probe_int_drift_for_model` into a model's ``(entry, dropped_columns)`` diff so the cascade fixed-point walk in @@ -1968,10 +1962,10 @@ def _merge_probe_drifts_into_diff( def _diff_one_sql_table_model( *, model: SlayerModel, - live_tables: Dict[str, "LiveTable"], - available_in_ds: Set[str], - probe_drifts: List[Tuple[str, DataType, DeleteReason]], -) -> Tuple[Optional[ToDeleteEntry], Set[str]]: + live_tables: dict[str, "LiveTable"], + available_in_ds: set[str], + probe_drifts: list[tuple[str, DataType, DeleteReason]], +) -> tuple[ToDeleteEntry | None, set[str]]: """Per-model body of :func:`_collect_sql_table_diffs` — resolves the live table, runs ``diff_sql_table_model``, and merges any DEV-1538 SQLite probe drifts so the cascade walk treats them as regular drops.""" @@ -1991,9 +1985,9 @@ def _diff_one_sql_table_model( async def _collect_sql_table_diffs( *, datasource: DatasourceConfig, - sql_table_models: List[SlayerModel], - available_in_ds: Set[str], -) -> Dict[str, Tuple[Optional[ToDeleteEntry], Set[str]]]: + sql_table_models: list[SlayerModel], + available_in_ds: set[str], +) -> dict[str, tuple[ToDeleteEntry | None, set[str]]]: """Run live SQLAlchemy introspection (off the event loop) and diff each sql_table-mode model against it. On SQLite, additionally run the DEV-1538 affinity probe per persisted INT base column and merge any @@ -2028,11 +2022,11 @@ async def _collect_sql_table_diffs( async def _collect_sql_diffs( *, datasource: DatasourceConfig, - sql_models: List[SlayerModel], - sql_clients: Optional[Dict[Tuple[str, str], SlayerSQLClient]], -) -> Dict[str, Tuple[Optional[ToDeleteEntry], Set[str]]]: + sql_models: list[SlayerModel], + sql_clients: dict[tuple[str, str], SlayerSQLClient] | None, +) -> dict[str, tuple[ToDeleteEntry | None, set[str]]]: """Trial-execute each sql-mode model concurrently and produce its diff.""" - out: Dict[str, Tuple[Optional[ToDeleteEntry], Set[str]]] = {} + out: dict[str, tuple[ToDeleteEntry | None, set[str]]] = {} if not sql_models: return out # DEV-1551: SlayerQueryEngine._sql_clients is tuple-keyed @@ -2056,9 +2050,9 @@ async def _diff_one(model: SlayerModel) -> None: async def validate_datasource( *, datasource: DatasourceConfig, - models: List[SlayerModel], - sql_clients: Optional[Dict[Tuple[str, str], SlayerSQLClient]] = None, -) -> List[ToDeleteEntry]: + models: list[SlayerModel], + sql_clients: dict[tuple[str, str], SlayerSQLClient] | None = None, +) -> list[ToDeleteEntry]: """Validate every persisted model in ``models`` (all in the same DS) against the live schema of ``datasource``. Read-only. """ diff --git a/slayer/facade/catalog.py b/slayer/facade/catalog.py index 84063464..556656f4 100644 --- a/slayer/facade/catalog.py +++ b/slayer/facade/catalog.py @@ -15,7 +15,6 @@ from __future__ import annotations import logging -from typing import Dict, FrozenSet, List, Optional, Set, Tuple from pydantic import BaseModel, ConfigDict, Field @@ -40,7 +39,7 @@ # expansion (§5.1 rule 3) and the custom-agg expansion (rule 4) both # skip these for built-ins. Custom aggs with non-empty ``params`` are # also skipped per rule 4 for the same reason. -_PARAMETRIC_BUILTIN_AGGS: FrozenSet[str] = frozenset({ +_PARAMETRIC_BUILTIN_AGGS: frozenset[str] = frozenset({ "weighted_avg", "percentile", "corr", "covar_samp", "covar_pop", }) @@ -50,16 +49,16 @@ class FacadeMetric(BaseModel): name: str - description: Optional[str] = None - label: Optional[str] = None - data_type: Optional[DataType] = None + description: str | None = None + label: str | None = None + data_type: DataType | None = None measure_formula: str class FacadeDimension(BaseModel): name: str - description: Optional[str] = None - label: Optional[str] = None + description: str | None = None + label: str | None = None data_type: DataType is_time: bool dimension_ref: str @@ -74,7 +73,7 @@ class FacadeJoin(BaseModel): exposed, mirroring the BFS dim/metric filter. """ target_model: str - join_pairs: List[List[str]] + join_pairs: list[list[str]] join_type: JoinType = JoinType.LEFT @@ -85,29 +84,29 @@ class FacadeTable(BaseModel): name: str table_type: str - description: Optional[str] = None - metrics: List[FacadeMetric] - dimensions: List[FacadeDimension] - joins: List[FacadeJoin] = Field(default_factory=list) + description: str | None = None + metrics: list[FacadeMetric] + dimensions: list[FacadeDimension] + joins: list[FacadeJoin] = Field(default_factory=list) # In-memory handle to the underlying SlayerModel — required by the # translator for ON-clause column validation (hidden FK/PK columns # don't appear on `dimensions`) and for the dynamic-join lookup # materialisation (DEV-1565). Excluded from any future serialisation # of FacadeCatalog. - model_ref: Optional[SlayerModel] = Field(default=None, exclude=True) + model_ref: SlayerModel | None = Field(default=None, exclude=True) class FacadeSchema(BaseModel): name: str - tables: List[FacadeTable] + tables: list[FacadeTable] class FacadeCatalog(BaseModel): catalog_name: str = CATALOG_NAME - schemas: List[FacadeSchema] + schemas: list[FacadeSchema] -def local_metrics(table: FacadeTable) -> List[FacadeMetric]: +def local_metrics(table: FacadeTable) -> list[FacadeMetric]: """Metrics that should appear in the per-table flat-column view (``pg_attribute`` / ``INFORMATION_SCHEMA.COLUMNS``) — saved ``ModelMeasure`` entries only. @@ -163,7 +162,7 @@ def local_metrics(table: FacadeTable) -> List[FacadeMetric]: ] -def local_dimensions(table: FacadeTable) -> List[FacadeDimension]: +def local_dimensions(table: FacadeTable) -> list[FacadeDimension]: """Mirror of :func:`local_metrics` for dimensions: drop cross-model entries (single-hop and multi-hop joined dimensions). Unlike metrics, dimensions don't have a synthetic-vs-user-authored split — @@ -174,7 +173,7 @@ def local_dimensions(table: FacadeTable) -> List[FacadeDimension]: def build_catalog( *, - models_by_datasource: Dict[str, List[SlayerModel]], + models_by_datasource: dict[str, list[SlayerModel]], bfs_depth: int = DEFAULT_BFS_DEPTH, ) -> FacadeCatalog: """Build a ``FacadeCatalog`` snapshot. @@ -184,10 +183,10 @@ def build_catalog( list_models(data_source=...)`` so cross-datasource joins are naturally constrained (SLayer doesn't auto-mirror joins across datasources). """ - schemas: List[FacadeSchema] = [] + schemas: list[FacadeSchema] = [] for datasource, models in models_by_datasource.items(): - by_name: Dict[str, SlayerModel] = {m.name: m for m in models} - tables: List[FacadeTable] = [] + by_name: dict[str, SlayerModel] = {m.name: m for m in models} + tables: list[FacadeTable] = [] for model in models: if model.hidden: continue @@ -231,7 +230,7 @@ def _column_types_supported(*, model: SlayerModel) -> bool: def _build_table( *, model: SlayerModel, - models_by_name: Dict[str, SlayerModel], + models_by_name: dict[str, SlayerModel], bfs_depth: int, ) -> FacadeTable: table_type = _table_type(model=model) @@ -253,13 +252,13 @@ def _build_table( def _facade_joins_for( - *, model: SlayerModel, models_by_name: Dict[str, SlayerModel], -) -> List[FacadeJoin]: + *, model: SlayerModel, models_by_name: dict[str, SlayerModel], +) -> list[FacadeJoin]: """Expose every direct (single-hop) join whose target is a non-hidden model in the same catalog. Mirrors the BFS filter so the translator's existence check never matches a join that isn't otherwise addressable (DEV-1565).""" - out: List[FacadeJoin] = [] + out: list[FacadeJoin] = [] for j in model.joins: target = models_by_name.get(j.target_model) if target is None or target.hidden: @@ -285,9 +284,9 @@ def _table_type(*, model: SlayerModel) -> str: def _walk_join_paths( *, root: SlayerModel, - models_by_name: Dict[str, SlayerModel], + models_by_name: dict[str, SlayerModel], max_depth: int, -) -> List[Tuple[List[str], SlayerModel]]: +) -> list[tuple[list[str], SlayerModel]]: """BFS the join graph from ``root`` up to ``max_depth`` hops. Returns a list of (path, target_model) tuples where ``path`` is the @@ -299,10 +298,10 @@ def _walk_join_paths( ``A→B→A`` revisit is allowed (a legitimate query shape when the join columns differ); past ``max_depth`` the BFS terminates. """ - out: List[Tuple[List[str], SlayerModel]] = [] + out: list[tuple[list[str], SlayerModel]] = [] if max_depth <= 0: return out - queue: List[Tuple[SlayerModel, List[str]]] = [(root, [])] + queue: list[tuple[SlayerModel, list[str]]] = [(root, [])] while queue: current, path = queue.pop(0) if len(path) >= max_depth: @@ -317,7 +316,7 @@ def _walk_join_paths( return out -def _path_dotted(path: List[str]) -> str: +def _path_dotted(path: list[str]) -> str: """Convert a join path to its dotted reference form. Used uniformly for both the catalog-facing metric / dimension ``name`` @@ -336,7 +335,7 @@ def _path_dotted(path: List[str]) -> str: # leaves BI tools (Metabase fingerprint, dbt schema scan) discovering # them as queryable columns whose execution then fails with # "Aggregation 'first' on measure '' requires a time column". -_TIME_DEPENDENT_AGGREGATIONS: FrozenSet[str] = frozenset({"first", "last"}) +_TIME_DEPENDENT_AGGREGATIONS: frozenset[str] = frozenset({"first", "last"}) def _model_has_resolvable_time_dimension(model: SlayerModel) -> bool: @@ -354,8 +353,8 @@ def _model_has_resolvable_time_dimension(model: SlayerModel) -> bool: def _eligible_aggregations( - *, column: Column, model: Optional[SlayerModel] = None, -) -> Set[str]: + *, column: Column, model: SlayerModel | None = None, +) -> set[str]: """Per §5.1.3: default-by-type ∩ explicit whitelist, with PK clamp. When ``model`` is given AND it has no time dimension, time-dependent @@ -375,7 +374,7 @@ def _eligible_aggregations( return base -def _eligible_custom_aggregations(*, model: SlayerModel) -> List[Aggregation]: +def _eligible_custom_aggregations(*, model: SlayerModel) -> list[Aggregation]: """Per §5.1.4: custom aggs that use only ``{value}`` (no extra params).""" return [agg for agg in model.aggregations if not agg.params] @@ -383,8 +382,8 @@ def _eligible_custom_aggregations(*, model: SlayerModel) -> List[Aggregation]: def _metric_expansion( *, model: SlayerModel, - reachable: List[Tuple[List[str], SlayerModel]], -) -> List[FacadeMetric]: + reachable: list[tuple[list[str], SlayerModel]], +) -> list[FacadeMetric]: local = _local_metrics_for(model=model) out = list(local) # Apply BFS-derived joined metrics. Rules 1-4 are computed on ``J`` @@ -433,7 +432,7 @@ def _synthetic_row_count(model: SlayerModel) -> FacadeMetric: ) -def _saved_model_measures(model: SlayerModel) -> List[FacadeMetric]: +def _saved_model_measures(model: SlayerModel) -> list[FacadeMetric]: """Rule 2: every saved ``ModelMeasure`` with a name.""" return [ FacadeMetric( @@ -448,7 +447,7 @@ def _saved_model_measures(model: SlayerModel) -> List[FacadeMetric]: ] -def _column_x_builtin_aggs(model: SlayerModel) -> List[FacadeMetric]: +def _column_x_builtin_aggs(model: SlayerModel) -> list[FacadeMetric]: """Rule 3: column × eligible-builtin-agg cartesian.""" return [ FacadeMetric( @@ -464,7 +463,7 @@ def _column_x_builtin_aggs(model: SlayerModel) -> List[FacadeMetric]: ] -def _column_x_custom_aggs(model: SlayerModel) -> List[FacadeMetric]: +def _column_x_custom_aggs(model: SlayerModel) -> list[FacadeMetric]: """Rule 4: column × parameterless custom aggs. Custom aggs are not gated by ``DEFAULT_AGGREGATIONS_BY_TYPE``, so we expose them on every non-hidden column. Custom-agg output type is opaque.""" @@ -485,7 +484,7 @@ def _column_x_custom_aggs(model: SlayerModel) -> List[FacadeMetric]: ] -def _local_metrics_for(*, model: SlayerModel) -> List[FacadeMetric]: +def _local_metrics_for(*, model: SlayerModel) -> list[FacadeMetric]: """Apply rules 1-4 to a single model in isolation (no join walk).""" return [ _synthetic_row_count(model), @@ -495,9 +494,9 @@ def _local_metrics_for(*, model: SlayerModel) -> List[FacadeMetric]: ] -def _local_dimensions_for(*, model: SlayerModel) -> List[FacadeDimension]: +def _local_dimensions_for(*, model: SlayerModel) -> list[FacadeDimension]: """Bare-column dims for a single model in isolation (no join walk).""" - out: List[FacadeDimension] = [] + out: list[FacadeDimension] = [] for col in model.columns: if col.hidden: continue @@ -514,7 +513,7 @@ def _local_dimensions_for(*, model: SlayerModel) -> List[FacadeDimension]: def build_local_view( model: SlayerModel, -) -> Tuple[List[FacadeDimension], List[FacadeMetric]]: +) -> tuple[list[FacadeDimension], list[FacadeMetric]]: """Build the bare-column dims + col×agg metrics for a single model in isolation (no join walk). Used by the translator's dynamic-join lookup materialisation (DEV-1565) so a join the catalog's BFS didn't @@ -524,13 +523,13 @@ def build_local_view( return _local_dimensions_for(model=model), _local_metrics_for(model=model) -def _describe_column_agg(*, column: Column, agg: str) -> Optional[str]: +def _describe_column_agg(*, column: Column, agg: str) -> str | None: if column.description: return f"{column.description} ({agg})" return None -def _agg_output_type(*, column: Column, agg: str) -> Optional[DataType]: +def _agg_output_type(*, column: Column, agg: str) -> DataType | None: """Coarse-grained output-type inference for column × agg pairs. Used only to populate ``INFORMATION_SCHEMA.METRICS.data_type``; the @@ -557,9 +556,9 @@ def _agg_output_type(*, column: Column, agg: str) -> Optional[DataType]: def _dimension_expansion( *, model: SlayerModel, - reachable: List[Tuple[List[str], SlayerModel]], -) -> List[FacadeDimension]: - out: List[FacadeDimension] = [] + reachable: list[tuple[list[str], SlayerModel]], +) -> list[FacadeDimension]: + out: list[FacadeDimension] = [] for col in model.columns: if col.hidden: continue diff --git a/slayer/facade/catalog_sql.py b/slayer/facade/catalog_sql.py index f338e36c..c1dea7f5 100644 --- a/slayer/facade/catalog_sql.py +++ b/slayer/facade/catalog_sql.py @@ -40,7 +40,7 @@ import json import logging import zlib -from typing import Any, Dict, List, Optional, Tuple +from typing import Any import duckdb import sqlglot @@ -70,7 +70,7 @@ # Well-known Postgres system catalog OIDs. Hardcoded so ``'pg_class'::regclass`` # resolves to ``1259`` (matching ``pg_description.classoid``) and Metabase's # get-tables JOIN works end-to-end. -KNOWN_SYSTEM_OIDS: Dict[str, int] = { +KNOWN_SYSTEM_OIDS: dict[str, int] = { "pg_class": 1259, "pg_namespace": 2615, "pg_attribute": 1249, @@ -92,7 +92,7 @@ OID_TIMESTAMP, ) -_TYPE_META: Dict[int, Tuple[str, int, str]] = { +_TYPE_META: dict[int, tuple[str, int, str]] = { OID_BOOL: ("bool", 1, "B"), OID_INT8: ("int8", 8, "N"), OID_TEXT: ("text", -1, "S"), @@ -104,12 +104,12 @@ # Inverse of _TYPE_META: type name → OID. Used to resolve ``::regtype`` # casts to the underlying ``pg_type.oid`` so catalog queries like # ``WHERE oid = 'int8'::regtype`` work. -_KNOWN_TYPE_OIDS: Dict[str, int] = { +_KNOWN_TYPE_OIDS: dict[str, int] = { typname: oid for oid, (typname, _len, _cat) in _TYPE_META.items() } -_UDT_NAME_BY_DATATYPE: Dict[DataType, str] = { +_UDT_NAME_BY_DATATYPE: dict[DataType, str] = { DataType.BOOLEAN: "bool", DataType.INT: "int8", DataType.TEXT: "text", @@ -138,16 +138,16 @@ class CatalogRelation(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) name: str - columns: List[FacadeColumn] - rows: List[Dict[str, Any]] + columns: list[FacadeColumn] + rows: list[dict[str, Any]] # --- corpus builder --------------------------------------------------------- def build_catalog_relations( - catalog: FacadeCatalog, datasource: Optional[str] = None, -) -> List[CatalogRelation]: + catalog: FacadeCatalog, datasource: str | None = None, +) -> list[CatalogRelation]: """Build every catalog table from ``catalog``. ``datasource`` is used as the ``catalog_name`` / ``table_catalog`` @@ -164,7 +164,7 @@ def build_catalog_relations( ``_is_*`` information_schema relations. """ ds = datasource or catalog.catalog_name - out: List[CatalogRelation] = [] + out: list[CatalogRelation] = [] out.append(_build_pg_namespace()) out.append(_build_pg_class(catalog)) out.append(_build_pg_attribute(catalog)) @@ -252,7 +252,7 @@ def _build_pg_class(catalog: FacadeCatalog) -> CatalogRelation: FacadeColumn(name="relispartition", type=DataType.BOOLEAN), ] rows = [] - seen_oids: Dict[int, str] = {} + seen_oids: dict[int, str] = {} for ds, tbl in _all_tables(catalog): oid = _table_oid(ds, tbl) _check_collision(seen_oids, oid, f"{ds}.{tbl.name}") @@ -381,7 +381,7 @@ def _build_pg_description(catalog: FacadeCatalog) -> CatalogRelation: FacadeColumn(name="objsubid", type=DataType.INT), FacadeColumn(name="description", type=DataType.TEXT), ] - rows: List[Dict[str, Any]] = [] + rows: list[dict[str, Any]] = [] for ds, tbl in _all_tables(catalog): oid = _table_oid(ds, tbl) if tbl.description: @@ -542,7 +542,7 @@ def _build_is_columns(catalog: FacadeCatalog, datasource: str) -> CatalogRelatio FacadeColumn(name="description", type=DataType.TEXT), FacadeColumn(name="label", type=DataType.TEXT), ] - rows: List[Dict[str, Any]] = [] + rows: list[dict[str, Any]] = [] # DEV-1567: exclude cross-model entries — they leak as dotted "columns" # that Metabase fingerprint scans then project (see local_metrics # docstring in slayer/facade/catalog.py). @@ -653,7 +653,7 @@ def _build_is_metrics( FacadeColumn(name="data_type", type=DataType.TEXT), FacadeColumn(name="label", type=DataType.TEXT), ] - rows: List[Dict[str, Any]] = [] + rows: list[dict[str, Any]] = [] for sch in catalog.schemas: for tbl in sch.tables: for m in tbl.metrics: @@ -686,7 +686,7 @@ def _build_is_dimensions( FacadeColumn(name="label", type=DataType.TEXT), FacadeColumn(name="is_time", type=DataType.BOOLEAN), ] - rows: List[Dict[str, Any]] = [] + rows: list[dict[str, Any]] = [] for sch in catalog.schemas: for tbl in sch.tables: for d in tbl.dimensions: @@ -700,7 +700,7 @@ def _build_is_dimensions( return CatalogRelation(name="_is_dimensions", columns=columns, rows=rows) -def _check_collision(seen: Dict[int, str], oid: int, key: str) -> None: +def _check_collision(seen: dict[int, str], oid: int, key: str) -> None: prior = seen.get(oid) if prior is not None and prior != key: raise ValueError( @@ -873,7 +873,7 @@ def _has_catalog_function_call(parsed: exp.Expression) -> bool: # Constant-return stubs are AST-rewritten directly to literals (no DuckDB # macros). This sidesteps DuckDB's lack of macro arity overloading — the # corpus has 2-arg and 3-arg variants of has_*_privilege. -_CONSTANT_STUB_LITERALS: Dict[str, Any] = { +_CONSTANT_STUB_LITERALS: dict[str, Any] = { "has_table_privilege": True, "has_any_column_privilege": True, "has_schema_privilege": True, @@ -887,7 +887,7 @@ def _has_catalog_function_call(parsed: exp.Expression) -> bool: # registered with a single arity per name. ``obj_description`` has two # Postgres arities; the 2-arg form is normalised to the 1-arg by dropping # the second argument at rewrite time. -_LOOKUP_STUB_NAMES: Dict[str, str] = { +_LOOKUP_STUB_NAMES: dict[str, str] = { "format_type": "_slayer_format_type", "obj_description": "_slayer_obj_description", "col_description": "_slayer_col_description", @@ -895,7 +895,7 @@ def _has_catalog_function_call(parsed: exp.Expression) -> bool: } -def _function_name_lower(node: exp.Expression) -> Optional[str]: +def _function_name_lower(node: exp.Expression) -> str | None: """Return the lowercased function name for any kind of function node, or None if ``node`` isn't a function call.""" if isinstance(node, exp.Anonymous): @@ -920,7 +920,7 @@ def _to_literal(value: Any) -> exp.Expression: return exp.Literal.string(str(value)) -def _unwrap_qualified_stub_call(node: exp.Expression) -> Optional[exp.Expression]: +def _unwrap_qualified_stub_call(node: exp.Expression) -> exp.Expression | None: """Detect ``information_schema.(args)`` (and the same for ``pg_catalog.``) and rewrite the qualified call as a bare private-name Anonymous so DuckDB resolves the macro. @@ -972,9 +972,9 @@ class _AstRewriter: """ def __init__(self, *, datasource: str, - regclass_map: Optional[Dict[str, int]] = None) -> None: + regclass_map: dict[str, int] | None = None) -> None: self.datasource = datasource - self.regclass_map: Dict[str, int] = regclass_map or {} + self.regclass_map: dict[str, int] = regclass_map or {} def rewrite(self, parsed: exp.Expression) -> exp.Expression: # Strip schema qualifiers AND rewrite information_schema names first @@ -1176,7 +1176,7 @@ def _rewrite_regclass_casts(self, node: exp.Expression) -> exp.Expression: return exp.Anonymous(this="slayer_regclass_oid", expressions=[inner]) @staticmethod - def _cast_target_kind(node: exp.Cast) -> Optional[str]: + def _cast_target_kind(node: exp.Cast) -> str | None: to = node.args.get("to") if to is None: return None @@ -1231,7 +1231,7 @@ def _substitute_context_functions(self, node: exp.Expression) -> exp.Expression: def _substitute_qualified_context_call( self, node: exp.Expression, - ) -> Optional[exp.Expression]: + ) -> exp.Expression | None: """Replace ``pg_catalog.`` (and ``pg_catalog.()``) as a whole so the outer ``Dot`` doesn't end up wrapping a string literal (``pg_catalog.'jaffle'`` — invalid SQL). @@ -1256,7 +1256,7 @@ def _substitute_qualified_context_call( def _substitute_qualified_context_column( self, node: exp.Expression, - ) -> Optional[exp.Expression]: + ) -> exp.Expression | None: """Replace ``pg_catalog.`` where sqlglot parses the whole thing as ``Column(this=, table='pg_catalog')`` — the no-parens shape (``pg_catalog.current_user``, @@ -1279,7 +1279,7 @@ def _substitute_qualified_context_column( return None return self._literal_for_context_name(str(ident.this).lower()) - def _substitute_dedicated_func(self, node: exp.Expression) -> Optional[exp.Expression]: + def _substitute_dedicated_func(self, node: exp.Expression) -> exp.Expression | None: """Dedicated sqlglot Func subclasses (typed nodes for niladic ctx fns).""" if isinstance(node, (exp.CurrentDatabase, getattr(exp, "CurrentCatalog", exp.CurrentDatabase))): return exp.Literal.string(self.datasource) @@ -1289,7 +1289,7 @@ def _substitute_dedicated_func(self, node: exp.Expression) -> Optional[exp.Expre return exp.Literal.string("public") return None - def _substitute_bareword_column(self, node: exp.Expression) -> Optional[exp.Expression]: + def _substitute_bareword_column(self, node: exp.Expression) -> exp.Expression | None: """sqlglot parses ``current_role`` (no parens) as a Column reference. Treat single-token unqualified Column refs naming a known niladic ctx function as that function. ``node.table`` is ``""`` (not None) for an @@ -1301,14 +1301,14 @@ def _substitute_bareword_column(self, node: exp.Expression) -> Optional[exp.Expr return None return self._literal_for_context_name(str(ident.this).lower()) - def _substitute_anonymous_function(self, node: exp.Expression) -> Optional[exp.Expression]: + def _substitute_anonymous_function(self, node: exp.Expression) -> exp.Expression | None: """Less-common Anonymous function spellings — fallback path.""" name = _function_name_lower(node) if name is None: return None return self._literal_for_context_name(name) - def _literal_for_context_name(self, name: str) -> Optional[exp.Expression]: + def _literal_for_context_name(self, name: str) -> exp.Expression | None: if name in {"current_database", "current_catalog", "currentdatabase", "currentcatalog"}: return exp.Literal.string(self.datasource) @@ -1405,7 +1405,7 @@ def _rewrite_regex_operators(node: exp.Expression) -> exp.Expression: # --- DuckDB type mapping ---------------------------------------------------- -_DUCKDB_TO_DATATYPE: Dict[str, DataType] = { +_DUCKDB_TO_DATATYPE: dict[str, DataType] = { # ints "TINYINT": DataType.INT, "SMALLINT": DataType.INT, "INTEGER": DataType.INT, "BIGINT": DataType.INT, "HUGEINT": DataType.INT, @@ -1442,7 +1442,7 @@ def _duckdb_typename_to_datatype(typename: str) -> DataType: # --- DuckDB type mapping for table columns at creation time ---------------- -_DATATYPE_TO_DUCKDB_CREATE: Dict[DataType, str] = { +_DATATYPE_TO_DUCKDB_CREATE: dict[DataType, str] = { DataType.INT: "BIGINT", DataType.DOUBLE: "DOUBLE", DataType.TEXT: "VARCHAR", @@ -1484,7 +1484,7 @@ def __init__(self, *, catalog: FacadeCatalog, datasource: str) -> None: pass # OID lookup for the regclass UDF: maps both schema-qualified and # bare names to OIDs (system catalogs + user tables). - self._regclass_map: Dict[str, int] = dict(KNOWN_SYSTEM_OIDS) + self._regclass_map: dict[str, int] = dict(KNOWN_SYSTEM_OIDS) for ds, tbl in _all_tables(catalog): oid = _table_oid(ds, tbl) self._regclass_map[f"public.{tbl.name}"] = oid @@ -1541,7 +1541,7 @@ def _register_stubs(self) -> None: ["VARCHAR"], "INTEGER", ) - def _regclass_oid(self, text: Optional[str]) -> int: + def _regclass_oid(self, text: str | None) -> int: if text is None: return 0 # Both schema-qualified and bare lookups go through the same map. @@ -1560,9 +1560,9 @@ def execute(self, *, parsed: exp.Expression, sql: str) -> RowBatch: raise TranslationError(str(exc)) from exc if description is None: return RowBatch(columns=[], rows=[]) - columns: List[FacadeColumn] = [] - col_keys: List[str] = [] - seen_keys: Dict[str, int] = {} + columns: list[FacadeColumn] = [] + col_keys: list[str] = [] + seen_keys: dict[str, int] = {} for col in description: name = col[0] typename = col[1] if len(col) > 1 else "VARCHAR" @@ -1631,7 +1631,7 @@ def _fingerprint(catalog: FacadeCatalog, datasource: str) -> str: def executor_for( - catalog: FacadeCatalog, datasource: Optional[str] = None, + catalog: FacadeCatalog, datasource: str | None = None, ) -> CatalogSqlExecutor: """Return a process-cached ``CatalogSqlExecutor`` for ``catalog``. diff --git a/slayer/facade/info_schema.py b/slayer/facade/info_schema.py index fbed9ef3..599865e1 100644 --- a/slayer/facade/info_schema.py +++ b/slayer/facade/info_schema.py @@ -23,7 +23,6 @@ from __future__ import annotations -from typing import List, Optional import sqlglot.expressions as exp @@ -49,7 +48,7 @@ _CATALOG_NAME_LOWER = CATALOG_NAME.lower() -def _is_information_schema_from(node: exp.Expression) -> Optional[str]: +def _is_information_schema_from(node: exp.Expression) -> str | None: """If ``node`` is ``SELECT ... FROM information_schema.
``, return the uppercased table name; else ``None``. @@ -94,7 +93,7 @@ def _is_information_schema_from(node: exp.Expression) -> Optional[str]: def match_info_schema( *, parsed: exp.Expression, catalog: FacadeCatalog, -) -> Optional[RowBatch]: +) -> RowBatch | None: """Return the canned ``INFORMATION_SCHEMA.
`` answer or ``None``.""" table_name = _is_information_schema_from(parsed) if table_name is None: @@ -126,7 +125,7 @@ def _serve_metrics(*, catalog: FacadeCatalog) -> RowBatch: FacadeColumn(name="data_type", type=DataType.TEXT), FacadeColumn(name="label", type=DataType.TEXT), ] - rows: List[dict] = [] + rows: list[dict] = [] for sch in catalog.schemas: for tbl in sch.tables: for m in tbl.metrics: @@ -153,7 +152,7 @@ def _serve_dimensions(*, catalog: FacadeCatalog) -> RowBatch: FacadeColumn(name="label", type=DataType.TEXT), FacadeColumn(name="is_time", type=DataType.BOOLEAN), ] - rows: List[dict] = [] + rows: list[dict] = [] for sch in catalog.schemas: for tbl in sch.tables: for d in tbl.dimensions: @@ -189,7 +188,7 @@ def _serve_tables(*, catalog: FacadeCatalog) -> RowBatch: FacadeColumn(name="table_name", type=DataType.TEXT), FacadeColumn(name="table_type", type=DataType.TEXT), ] - rows: List[dict] = [] + rows: list[dict] = [] for sch in catalog.schemas: for tbl in sch.tables: rows.append({ @@ -224,7 +223,7 @@ def _serve_columns(*, catalog: FacadeCatalog) -> RowBatch: FacadeColumn(name="is_nullable", type=DataType.TEXT), # Postgres YES/NO FacadeColumn(name="column_kind", type=DataType.TEXT), # METRIC / DIMENSION ] - rows: List[dict] = [] + rows: list[dict] = [] for sch in catalog.schemas: for tbl in sch.tables: position = 1 diff --git a/slayer/facade/probe_queries.py b/slayer/facade/probe_queries.py index a551a03b..d3c394bc 100644 --- a/slayer/facade/probe_queries.py +++ b/slayer/facade/probe_queries.py @@ -13,7 +13,6 @@ from __future__ import annotations -from typing import Optional import sqlglot.expressions as exp @@ -135,7 +134,7 @@ def _matches_select_current_database(node: exp.Expression) -> bool: return False -def match_probe(parsed: exp.Expression) -> Optional[RowBatch]: +def match_probe(parsed: exp.Expression) -> RowBatch | None: """Return the canned ``RowBatch`` for a matching probe, else ``None``.""" if _matches_select_one(parsed): return _batch_select_one() diff --git a/slayer/facade/rows.py b/slayer/facade/rows.py index 6d297f58..b6bf2494 100644 --- a/slayer/facade/rows.py +++ b/slayer/facade/rows.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Any, Dict, List +from typing import Any from pydantic import BaseModel, ConfigDict @@ -26,5 +26,5 @@ class FacadeColumn(BaseModel): class RowBatch(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - columns: List[FacadeColumn] - rows: List[Dict[str, Any]] + columns: list[FacadeColumn] + rows: list[dict[str, Any]] diff --git a/slayer/facade/translator.py b/slayer/facade/translator.py index 3e83c14e..6c97673e 100644 --- a/slayer/facade/translator.py +++ b/slayer/facade/translator.py @@ -37,7 +37,7 @@ import logging import re from contextvars import ContextVar -from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union +from collections.abc import Callable, Sequence import sqlglot import sqlglot.errors @@ -104,7 +104,7 @@ class SetSettingOp(BaseModel): class ResetSettingOp(BaseModel): """``RESET `` (``name`` set) or ``RESET ALL`` (``reset_all=True``).""" - name: Optional[str] = None + name: str | None = None reset_all: bool = False @@ -121,11 +121,11 @@ class ProbeMatcherOutcome(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) batch: RowBatch - settings_mutation: Optional[SetSettingOp] = None + settings_mutation: SetSettingOp | None = None ProbeMatcher = Callable[ - [exp.Expression], Optional[Union[RowBatch, ProbeMatcherOutcome]], + [exp.Expression], RowBatch | ProbeMatcherOutcome | None, ] @@ -147,7 +147,7 @@ class ProbeResult(TranslatorResult): matchers return ``None`` here.""" batch: RowBatch - settings_mutation: Optional[SetSettingOp] = None + settings_mutation: SetSettingOp | None = None class InfoSchemaResult(TranslatorResult): @@ -177,9 +177,9 @@ class NoOpResult(TranslatorResult): intent. Both are PG-facade-only; the Flight facade ignores them. """ - command_tag: Optional[str] = None - set_setting: Optional[SetSettingOp] = None - reset_setting: Optional[ResetSettingOp] = None + command_tag: str | None = None + set_setting: SetSettingOp | None = None + reset_setting: ResetSettingOp | None = None class QueryResult(TranslatorResult): @@ -198,10 +198,10 @@ class QueryResult(TranslatorResult): """ query: SlayerQuery - column_name_mapping: List[Tuple[str, str]] + column_name_mapping: list[tuple[str, str]] facade_table: FacadeTable schema_name: str - projection_types: "List[Optional['DataType']]" + projection_types: "list[DataType | None]" @property def flight_table(self) -> FacadeTable: @@ -240,7 +240,7 @@ def __init__(self, message: str, *, status: str = "INVALID_ARGUMENT") -> None: # --- AST helpers ------------------------------------------------------------- -_TIME_GRAIN_NAMES: Dict[str, TimeGranularity] = { +_TIME_GRAIN_NAMES: dict[str, TimeGranularity] = { "year": TimeGranularity.YEAR, "quarter": TimeGranularity.QUARTER, "month": TimeGranularity.MONTH, @@ -253,7 +253,7 @@ def __init__(self, message: str, *, status: str = "INVALID_ARGUMENT") -> None: # sqlglot represents the unwrapped one-arg time functions as dedicated nodes # (exp.Month, exp.Year, …). date_trunc is exp.DateTrunc with a literal unit. -_TIME_GRAIN_CLASSES: Dict[type, TimeGranularity] = { +_TIME_GRAIN_CLASSES: dict[type, TimeGranularity] = { exp.Year: TimeGranularity.YEAR, exp.Quarter: TimeGranularity.QUARTER, exp.Month: TimeGranularity.MONTH, @@ -265,14 +265,14 @@ def __init__(self, message: str, *, status: str = "INVALID_ARGUMENT") -> None: # sqlglot aggregate-function AST classes → SLayer aggregation names. COUNT is # handled separately (it has *-arg and DISTINCT variants). -_AGG_CLASS_TO_NAME: Dict[type, str] = { +_AGG_CLASS_TO_NAME: dict[type, str] = { exp.Sum: "sum", exp.Avg: "avg", exp.Min: "min", exp.Max: "max", } -_COMPARATOR_SQL: Dict[type, str] = { +_COMPARATOR_SQL: dict[type, str] = { exp.GT: ">", exp.GTE: ">=", exp.LT: "<", @@ -304,7 +304,7 @@ def __init__(self, message: str, *, status: str = "INVALID_ARGUMENT") -> None: # Callers needing exact NUMERIC precision, narrow integer wire widths, or # TZ-aware timestamps must compute upstream. See pg-facade.md # §"CAST coarse-OID mapping" for the user-facing table. -_SQLGLOT_TYPE_TO_DATATYPE: Dict[exp.DataType.Type, DataType] = { +_SQLGLOT_TYPE_TO_DATATYPE: dict[exp.DataType.Type, DataType] = { exp.DataType.Type.TEXT: DataType.TEXT, exp.DataType.Type.VARCHAR: DataType.TEXT, exp.DataType.Type.CHAR: DataType.TEXT, @@ -332,7 +332,7 @@ def __init__(self, message: str, *, status: str = "INVALID_ARGUMENT") -> None: # other target we'd push the failure into value_to_text / value_to_binary at # response time, surfacing as an opaque connection error instead of a clean # translation error. Identity (X→X) is always admitted; computed implicitly. -_CAST_ADMITTED_TARGETS: Dict[DataType, frozenset] = { +_CAST_ADMITTED_TARGETS: dict[DataType, frozenset] = { DataType.DATE: frozenset({DataType.DATE, DataType.TIMESTAMP, DataType.TEXT}), DataType.TIMESTAMP: frozenset({DataType.TIMESTAMP, DataType.DATE, DataType.TEXT}), DataType.INT: frozenset({DataType.INT, DataType.DOUBLE, DataType.TEXT}), @@ -379,7 +379,7 @@ def __init__(self, message: str, *, status: str = "INVALID_ARGUMENT") -> None: _GROUP_BY_KIND = "GROUP BY" -def _sqlglot_type_to_datatype(node: exp.DataType) -> Optional[DataType]: +def _sqlglot_type_to_datatype(node: exp.DataType) -> DataType | None: """Map a sqlglot ``DataType`` node to a SLayer ``DataType``. Returns ``None`` when the type member isn't in the mapping (UUID, JSON, @@ -389,7 +389,7 @@ def _sqlglot_type_to_datatype(node: exp.DataType) -> Optional[DataType]: return _SQLGLOT_TYPE_TO_DATATYPE.get(node.this) -def _is_admitted_cast(source: Optional[DataType], target: DataType) -> bool: +def _is_admitted_cast(source: DataType | None, target: DataType) -> bool: """Strict per-pair admission. ``None`` source admits only ``TEXT``.""" if source is None: return target is DataType.TEXT @@ -401,8 +401,8 @@ def _is_admitted_cast(source: Optional[DataType], target: DataType) -> bool: def _column_to_dotted( col: exp.Column, *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> str: """Reconstruct the dotted reference from a sqlglot ``Column``. @@ -423,7 +423,7 @@ def _column_to_dotted( built — so ``"Stores"."name"`` with ``alias_map={"Stores": "stores"}`` yields ``"stores.name"`` (the SLayer cross-model dotted form). """ - parts: List[str] = [] + parts: list[str] = [] for key in ("catalog", "db", "table"): node = col.args.get(key) if node is None: @@ -436,8 +436,8 @@ def _column_to_dotted( def _apply_alias_remap( - parts: List[str], alias_map: Optional[Dict[str, str]], -) -> List[str]: + parts: list[str], alias_map: dict[str, str] | None, +) -> list[str]: """If ``parts`` is at least 2-part and the table-qualifier (second-to- last element) matches an entry in ``alias_map`` (case-insensitive), rewrite it to the target model name. Otherwise return ``parts`` @@ -465,8 +465,8 @@ def _apply_alias_remap( def _apply_strip_prefix( - parts: List[str], strip_prefix: Optional[Tuple[str, str]], -) -> List[str]: + parts: list[str], strip_prefix: tuple[str, str] | None, +) -> list[str]: """Drop the leading ``schema.table.`` qualifier from ``parts`` when it matches ``strip_prefix``. Four-part catalog-qualified refs drop the leading 3; three-part drops the leading 2; two-part drops the @@ -499,7 +499,7 @@ def _apply_strip_prefix( def _detect_time_grain_date_trunc( node: exp.Expression, -) -> Optional[Tuple[TimeGranularity, exp.Column]]: +) -> tuple[TimeGranularity, exp.Column] | None: """Plain ``DATE_TRUNC(, )`` detector. Does NOT unwrap any day offsets on the column side — a bare @@ -530,7 +530,7 @@ def _detect_time_grain_date_trunc( return grain, col -def _date_trunc_unit(node: exp.Expression) -> Optional[str]: +def _date_trunc_unit(node: exp.Expression) -> str | None: """Return the lowercase unit string of a ``DATE_TRUNC``/``TIMESTAMP_TRUNC`` node, or ``None`` if not a trunc call. Used by the Sunday-week wrapper detector below.""" @@ -546,7 +546,7 @@ def _date_trunc_unit(node: exp.Expression) -> Optional[str]: def _detect_sunday_week_wrapper( node: exp.Expression, -) -> Optional[Tuple[TimeGranularity, exp.Column]]: +) -> tuple[TimeGranularity, exp.Column] | None: """Recognise the complete Metabase Sunday-week wrapper as a single shape. Full pattern (the one Metabase emits for a week breakout on a DATE @@ -590,7 +590,7 @@ def _detect_sunday_week_wrapper( return TimeGranularity.WEEK_SUNDAY, unwrapped_col -def _day_interval_sign(node: exp.Expression) -> Optional[int]: +def _day_interval_sign(node: exp.Expression) -> int | None: """Return ``+1`` for ``INTERVAL '1 day'``, ``-1`` for ``INTERVAL '-1 day'``, or ``None`` if ``node`` isn't a one-day interval at all. @@ -659,7 +659,7 @@ def _unwrap_signed_day_offset( def _detect_time_grain_single_arg( node: exp.Expression, -) -> Optional[Tuple[TimeGranularity, exp.Column]]: +) -> tuple[TimeGranularity, exp.Column] | None: """Dedicated AST classes like ``exp.Month`` / ``exp.Year``.""" for cls, grain in _TIME_GRAIN_CLASSES.items(): if isinstance(node, cls): @@ -672,7 +672,7 @@ def _detect_time_grain_single_arg( def _detect_time_grain_anonymous( node: exp.Anonymous, -) -> Optional[Tuple[TimeGranularity, exp.Column]]: +) -> tuple[TimeGranularity, exp.Column] | None: """``hour(col)`` / ``minute(col)`` / ``second(col)`` come through here.""" grain = _TIME_GRAIN_NAMES.get(str(node.this).lower()) if grain is None: @@ -683,7 +683,7 @@ def _detect_time_grain_anonymous( return None -def _detect_time_grain(node: exp.Expression) -> Optional[Tuple[TimeGranularity, exp.Column]]: +def _detect_time_grain(node: exp.Expression) -> tuple[TimeGranularity, exp.Column] | None: """If ``node`` is ``()`` or ``date_trunc('', )``, return ``(granularity, column)``. Otherwise ``None``. @@ -727,8 +727,8 @@ def _detect_time_grain(node: exp.Expression) -> Optional[Tuple[TimeGranularity, def _alias_for_time_grain( grain: TimeGranularity, col: exp.Column, *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> str: """The flat projection name we expose for ``month(ordered_at)`` etc. @@ -747,7 +747,7 @@ class _AggCall(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) agg: str # SLayer aggregation name: sum/avg/min/max/count/count_distinct - inner_ref: Optional[str] = None # dotted column ref, or None for COUNT(*) + inner_ref: str | None = None # dotted column ref, or None for COUNT(*) inner_is_column: bool = False # False → COUNT(*) or non-column arg is_count_star: bool = False # True only for the literal COUNT(*) @@ -755,9 +755,9 @@ class _AggCall(BaseModel): def _detect_aggregate( # NOSONAR(S3776) — flat per-aggregate-kind dispatch; splitting hides the shape node: exp.Expression, *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, -) -> Optional[_AggCall]: + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, +) -> _AggCall | None: """If ``node`` is a SQL aggregate call, classify it; else ``None``. ``COUNT(*)`` → ``count`` / ``inner_ref=None``. ``COUNT(DISTINCT col)`` → @@ -846,7 +846,7 @@ def _assert_local_metric(metric: FacadeMetric) -> None: def _metric_for_aggregate( - call: _AggCall, table: FacadeTable, metrics_by_formula: Dict[str, FacadeMetric], + call: _AggCall, table: FacadeTable, metrics_by_formula: dict[str, FacadeMetric], ) -> FacadeMetric: """Resolve an aggregate call to its catalog ``FacadeMetric``. @@ -878,16 +878,16 @@ def _metric_for_aggregate( # --- table resolution -------------------------------------------------------- -def _flatten_catalog(catalog: FacadeCatalog) -> Dict[str, List[Tuple[str, FacadeTable]]]: +def _flatten_catalog(catalog: FacadeCatalog) -> dict[str, list[tuple[str, FacadeTable]]]: """Build a (model_name → [(schema, table), …]) index for bare-name lookup.""" - by_name: Dict[str, List[Tuple[str, FacadeTable]]] = {} + by_name: dict[str, list[tuple[str, FacadeTable]]] = {} for sch in catalog.schemas: for tbl in sch.tables: by_name.setdefault(tbl.name, []).append((sch.name, tbl)) return by_name -def _unwrap_identifier(node: Optional[exp.Expression]) -> Optional[str]: +def _unwrap_identifier(node: exp.Expression | None) -> str | None: """Pull the string value out of a sqlglot identifier-ish node.""" if node is None: return None @@ -896,7 +896,7 @@ def _unwrap_identifier(node: Optional[exp.Expression]) -> Optional[str]: def _resolve_qualified_table( *, schema_str: str, table_name: str, catalog: FacadeCatalog, -) -> Tuple[str, FacadeTable]: +) -> tuple[str, FacadeTable]: # Try the exact schema match first — if a catalog actually carries a # ``public`` schema (or whatever name the user passed), honour the # explicit qualifier. @@ -921,7 +921,7 @@ def _resolve_qualified_table( def _resolve_bare_table( *, table_name: str, catalog: FacadeCatalog, -) -> Tuple[str, FacadeTable]: +) -> tuple[str, FacadeTable]: matches = _flatten_catalog(catalog).get(table_name, []) if not matches: raise TranslationError(f"Unknown table: {table_name!r}") @@ -936,7 +936,7 @@ def _resolve_bare_table( def _resolve_table( from_clause: exp.From, catalog: FacadeCatalog, -) -> Tuple[str, FacadeTable]: +) -> tuple[str, FacadeTable]: """Resolve a SELECT's FROM into ``(schema_name, FacadeTable)``. Handles the three qualification forms (§6.1): @@ -979,25 +979,25 @@ class _ProjectionItem(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) projected_name: str # what the BI tool sees (alias or natural name) - metric: Optional[FacadeMetric] = None - dimension: Optional[FacadeDimension] = None - time_grain: Optional[TimeGranularity] = None - time_grain_underlying: Optional[FacadeDimension] = None + metric: FacadeMetric | None = None + dimension: FacadeDimension | None = None + time_grain: TimeGranularity | None = None + time_grain_underlying: FacadeDimension | None = None # DEV-1566: target type when the projection was CAST( AS ). # Overrides projection_types[i] at _record_metric / _record_dimension time, # leaving engine_alias and the SlayerQuery measure/dimension untouched. - cast_target: Optional[DataType] = None + cast_target: DataType | None = None def _resolve_time_grain_projection( *, grain: TimeGranularity, col: exp.Column, - alias_name: Optional[str], + alias_name: str | None, table: FacadeTable, - dims_by_name: Dict[str, FacadeDimension], - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + dims_by_name: dict[str, FacadeDimension], + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> _ProjectionItem: dotted = _column_to_dotted(col, strip_prefix=strip_prefix, alias_map=alias_map) dim = dims_by_name.get(dotted) @@ -1024,9 +1024,9 @@ def _resolve_time_grain_projection( def _resolve_aggregate_projection( *, call: _AggCall, - alias_name: Optional[str], + alias_name: str | None, table: FacadeTable, - metrics_by_formula: Dict[str, FacadeMetric], + metrics_by_formula: dict[str, FacadeMetric], ) -> _ProjectionItem: """Map a SQL aggregate call to the same projection item a bare metric name would produce (DEV-1486 decision 21).""" @@ -1042,12 +1042,12 @@ def _resolve_aggregate_projection( def _resolve_column_projection( *, body: exp.Column, - alias_name: Optional[str], + alias_name: str | None, table: FacadeTable, - metrics_by_name: Dict[str, FacadeMetric], - dims_by_name: Dict[str, FacadeDimension], - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + metrics_by_name: dict[str, FacadeMetric], + dims_by_name: dict[str, FacadeDimension], + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> _ProjectionItem: dotted = _column_to_dotted(body, strip_prefix=strip_prefix, alias_map=alias_map) if dotted in metrics_by_name: @@ -1069,7 +1069,7 @@ def _resolve_column_projection( def _detect_column_cast( body: exp.Expression, -) -> Optional[Tuple[exp.Column, DataType]]: +) -> tuple[exp.Column, DataType] | None: """If ``body`` is ``CAST( AS )``, return ``(inner_col, target_data_type)``. Otherwise ``None``. @@ -1092,12 +1092,12 @@ def _resolve_column_cast_projection( *, body: exp.Column, cast_target: DataType, - alias_name: Optional[str], + alias_name: str | None, table: FacadeTable, - metrics_by_name: Dict[str, FacadeMetric], - dims_by_name: Dict[str, FacadeDimension], - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + metrics_by_name: dict[str, FacadeMetric], + dims_by_name: dict[str, FacadeDimension], + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> _ProjectionItem: """Resolve ``CAST( AS )`` to a projection item that runs the bare column through the engine and overrides the wire OID via @@ -1126,7 +1126,7 @@ def _resolve_column_cast_projection( # DEV-1558 B5: hygiene-scalar wrappers Metabase uses for fingerprint queries. # Each maps a sqlglot AST class to a human-readable name for the WARNING log. -_HYGIENE_FUNC_CLASSES: Dict[type, str] = { +_HYGIENE_FUNC_CLASSES: dict[type, str] = { exp.Substring: "SUBSTRING", exp.Upper: "UPPER", exp.Lower: "LOWER", @@ -1145,7 +1145,7 @@ def _resolve_column_cast_projection( "upper", "lower", "trim", "length"} -def _detect_hygiene_wrapper(body: exp.Expression) -> Optional[Tuple[str, exp.Column]]: +def _detect_hygiene_wrapper(body: exp.Expression) -> tuple[str, exp.Column] | None: """If ``body`` is a hygiene-scalar wrapper around exactly one column reference, return ``(printable_func_name, inner_col)``. Otherwise None. """ @@ -1165,7 +1165,7 @@ def _detect_hygiene_wrapper(body: exp.Expression) -> Optional[Tuple[str, exp.Col def _is_fingerprint_shape_wrap( - inner_col: exp.Column, *, strip_prefix: Optional[Tuple[str, str]], + inner_col: exp.Column, *, strip_prefix: tuple[str, str] | None, ) -> bool: """True iff ``inner_col`` is shaped like Metabase's fingerprint projection — a 3-part qualified column reference whose @@ -1192,9 +1192,9 @@ def _is_fingerprint_shape_wrap( return len(stripped) < len(raw) -def _raw_column_parts(col: exp.Column) -> List[str]: +def _raw_column_parts(col: exp.Column) -> list[str]: """The qualifier parts plus leaf identifier of ``col``, in order.""" - parts: List[str] = [] + parts: list[str] = [] for key in ("catalog", "db", "table"): node = col.args.get(key) if node is None: @@ -1208,10 +1208,10 @@ def _raw_column_parts(col: exp.Column) -> List[str]: def _build_projection_lookups( table: FacadeTable, *, - extra_dims_by_name: Optional[Dict[str, FacadeDimension]] = None, - extra_metrics_by_name: Optional[Dict[str, FacadeMetric]] = None, - extra_metrics_by_formula: Optional[Dict[str, FacadeMetric]] = None, -) -> Tuple[Dict[str, FacadeMetric], Dict[str, FacadeMetric], Dict[str, FacadeDimension]]: + extra_dims_by_name: dict[str, FacadeDimension] | None = None, + extra_metrics_by_name: dict[str, FacadeMetric] | None = None, + extra_metrics_by_formula: dict[str, FacadeMetric] | None = None, +) -> tuple[dict[str, FacadeMetric], dict[str, FacadeMetric], dict[str, FacadeDimension]]: """Assemble the (metrics_by_name, metrics_by_formula, dims_by_name) lookup dicts from the catalog's table metadata, overlaid with the DEV-1565 dynamic-join extras when supplied.""" @@ -1230,13 +1230,13 @@ def _build_projection_lookups( def _resolve_projection( # NOSONAR(S3776) — flat dispatch over projection-expression shapes; one branch per shape by design expressions: Sequence[exp.Expression], table: FacadeTable, *, - schema_name: Optional[str] = None, - alias_map: Optional[Dict[str, str]] = None, - extra_dims_by_name: Optional[Dict[str, FacadeDimension]] = None, - extra_metrics_by_name: Optional[Dict[str, FacadeMetric]] = None, - extra_metrics_by_formula: Optional[Dict[str, FacadeMetric]] = None, + schema_name: str | None = None, + alias_map: dict[str, str] | None = None, + extra_dims_by_name: dict[str, FacadeDimension] | None = None, + extra_metrics_by_name: dict[str, FacadeMetric] | None = None, + extra_metrics_by_formula: dict[str, FacadeMetric] | None = None, allow_column_cast: bool = True, -) -> List[_ProjectionItem]: +) -> list[_ProjectionItem]: """Walk the projection list, classifying each item against the table. DEV-1565: when a LEFT JOIN to a dynamic-fallback target is in play, @@ -1260,16 +1260,16 @@ def _resolve_projection( # NOSONAR(S3776) — flat dispatch over projection-exp # DEV-1558 B5: when the FROM was schema-qualified (e.g. "public"."orders"), # let column refs that lead with the same `schema.table.` prefix resolve # to the bare column. - strip_prefix: Optional[Tuple[str, str]] = ( + strip_prefix: tuple[str, str] | None = ( (schema_name, table.name) if schema_name else None ) - out: List[_ProjectionItem] = [] + out: list[_ProjectionItem] = [] for expr in expressions: if isinstance(expr, exp.Star): raise TranslationError(SELECT_STAR_MESSAGE) - alias_name: Optional[str] = None + alias_name: str | None = None body: exp.Expression = expr if isinstance(expr, exp.Alias): alias_name = str(expr.alias) @@ -1350,9 +1350,9 @@ def _resolve_projection( # NOSONAR(S3776) — flat dispatch over projection-exp # --- WHERE translation ------------------------------------------------------- -def _split_and_chain(node: exp.Expression) -> List[exp.Expression]: +def _split_and_chain(node: exp.Expression) -> list[exp.Expression]: """Flatten a top-level AND chain into its conjuncts.""" - out: List[exp.Expression] = [] + out: list[exp.Expression] = [] stack = [node] while stack: cur = stack.pop() @@ -1367,9 +1367,9 @@ def _split_and_chain(node: exp.Expression) -> List[exp.Expression]: def _lift_time_between( conj: exp.Between, time_dim_names: set[str], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, -) -> Optional[Tuple[str, Optional[str], Optional[str]]]: + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, +) -> tuple[str, str | None, str | None] | None: col = conj.this if not isinstance(col, exp.Column): return None @@ -1386,9 +1386,9 @@ def _lift_time_between( def _lift_time_comparator( conj: exp.Expression, time_dim_names: set[str], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, -) -> Optional[Tuple[str, Optional[str], Optional[str]]]: + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, +) -> tuple[str, str | None, str | None] | None: col = conj.this if not isinstance(col, exp.Column): return None @@ -1405,9 +1405,9 @@ def _lift_time_comparator( def _aggregate_alias_for_column( side: exp.Expression, - items_by_projected_name: Dict[str, "_ProjectionItem"], - *, strip_prefix: Optional[Tuple[str, str]] = None, -) -> Optional["_ProjectionItem"]: + items_by_projected_name: dict[str, "_ProjectionItem"], + *, strip_prefix: tuple[str, str] | None = None, +) -> "_ProjectionItem | None": """If ``side`` is a column reference whose dotted (strip_prefix-aware) name matches an aggregate projection's ``projected_name``, return that projection item; else ``None``. @@ -1429,9 +1429,9 @@ def _aggregate_alias_for_column( def _try_aggregate_alias_filter( conj: exp.Expression, - items_by_projected_name: Dict[str, "_ProjectionItem"], - *, strip_prefix: Optional[Tuple[str, str]] = None, -) -> Optional[str]: + items_by_projected_name: dict[str, "_ProjectionItem"], + *, strip_prefix: tuple[str, str] | None = None, +) -> str | None: """If ``conj`` is `` `` (in either order), return the colon-form filter string. If one side matches an aggregate alias but the other side is not a literal, raise. Return @@ -1466,9 +1466,9 @@ def _try_aggregate_alias_filter( def _classify_where_conjunct( conj: exp.Expression, time_dim_names: set[str], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, -) -> Tuple[Optional[Tuple[str, Optional[str], Optional[str]]], Optional[str]]: + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, +) -> tuple[tuple[str, str | None, str | None] | None, str | None]: """Classify a single conjunct. Returns ``((time_dim, date_range_lo, date_range_hi), None)`` if this is @@ -1504,8 +1504,8 @@ def _classify_where_conjunct( def _normalise_predicate_columns( node: exp.Expression, *, - strip_prefix: Optional[Tuple[str, str]], - alias_map: Optional[Dict[str, str]] = None, + strip_prefix: tuple[str, str] | None, + alias_map: dict[str, str] | None = None, ) -> exp.Expression: """Walk ``node`` and rewrite every ``exp.Column``: apply ``alias_map`` first (so ``.`` becomes ``.``), then @@ -1537,7 +1537,7 @@ def rewrite(child: exp.Expression) -> exp.Expression: return node.transform(rewrite) -def _literal_str(node: Optional[exp.Expression]) -> Optional[str]: +def _literal_str(node: exp.Expression | None) -> str | None: if node is None: return None if isinstance(node, exp.Literal): @@ -1568,7 +1568,7 @@ def _quote_keyword_aliases(sql: str) -> str: def _parse_with_keyword_alias_fallback( - sql: str, *, dialect: Optional[str], + sql: str, *, dialect: str | None, ) -> exp.Expression: """Parse ``sql`` with sqlglot. On failure, retry once with unquoted SQL keyword aliases auto-quoted — Metabase's table-privileges CTE @@ -1588,13 +1588,13 @@ def _parse_with_keyword_alias_fallback( def _apply_where( - where: Optional[exp.Where], - time_dims_built: Dict[str, TimeDimension], - items_by_projected_name: Dict[str, "_ProjectionItem"], - filters_out: List[str], + where: exp.Where | None, + time_dims_built: dict[str, TimeDimension], + items_by_projected_name: dict[str, "_ProjectionItem"], + filters_out: list[str], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> None: """Walk the WHERE chain; route aggregate-alias refs to colon-form filters (DEV-1568), lift time-dim filters, append verbatim rest.""" @@ -1625,14 +1625,14 @@ def _apply_where( def _apply_having( - having: Optional[exp.Having], + having: exp.Having | None, table: FacadeTable, - items_by_projected_name: Dict[str, "_ProjectionItem"], - filters_out: List[str], + items_by_projected_name: dict[str, "_ProjectionItem"], + filters_out: list[str], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, - extra_metrics_by_formula: Optional[Dict[str, FacadeMetric]] = None, + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, + extra_metrics_by_formula: dict[str, FacadeMetric] | None = None, ) -> None: """Map ``HAVING `` (DEV-1486 decision 21) and ``HAVING `` (DEV-1568) conjuncts to @@ -1687,7 +1687,7 @@ def _flip_comparator(op_sql: str) -> str: # --- ORDER BY / GROUP BY ----------------------------------------------------- -def _item_cast_source_type(item: _ProjectionItem) -> Optional[DataType]: +def _item_cast_source_type(item: _ProjectionItem) -> DataType | None: """Return the underlying source DataType of a CAST-projected item, or None when the source type is unknown (custom metric without a declared ``data_type``). Caller MUST gate on ``item.cast_target is not None`` @@ -1702,7 +1702,7 @@ def _item_cast_source_type(item: _ProjectionItem) -> Optional[DataType]: def _cast_pair_is_lossy( - *, source: Optional[DataType], target: DataType, + *, source: DataType | None, target: DataType, lossy_pairs: frozenset, ) -> bool: """A cast is lossy when either: (a) the source is known and (source, @@ -1718,7 +1718,7 @@ def _cast_pair_is_lossy( def _lossy_cast_error_message( - *, kind: str, name: str, source: Optional[DataType], target: DataType, + *, kind: str, name: str, source: DataType | None, target: DataType, ) -> str: """One-line user-facing message for the lossy-CAST-pair rejection. The ``kind`` is ``ORDER BY`` or ``GROUP BY``; ``name`` is the user-visible @@ -1757,12 +1757,12 @@ def _reject_lossy_cast_or_pass( def _resolve_order_by_item( body: exp.Expression, - item_by_projected_name: Dict[str, _ProjectionItem], - metric_item_by_formula: Dict[str, _ProjectionItem], + item_by_projected_name: dict[str, _ProjectionItem], + metric_item_by_formula: dict[str, _ProjectionItem], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, -) -> Tuple[_ProjectionItem, str]: + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, +) -> tuple[_ProjectionItem, str]: """Look up the projection item an ORDER BY term resolves to. Tries the bare-name lookup first (alias or canonical metric name from @@ -1788,12 +1788,12 @@ def _resolve_order_by_item( def _translate_order_by( - order: Optional[exp.Order], - item_by_projected_name: Dict[str, _ProjectionItem], + order: exp.Order | None, + item_by_projected_name: dict[str, _ProjectionItem], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, -) -> List[OrderItem]: + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, +) -> list[OrderItem]: if order is None: return [] # DEV-1568: when a SELECT aliases an aggregate (``SUM(revenue) AS "rev"``) @@ -1802,12 +1802,12 @@ def _translate_order_by( # — but ``item_by_projected_name`` is keyed on the user alias (``"rev"``). # Build a fallback map keyed by ``metric.measure_formula`` so a repeated # aggregate call resolves to the same projection the alias does. - metric_item_by_formula: Dict[str, _ProjectionItem] = { + metric_item_by_formula: dict[str, _ProjectionItem] = { item.metric.measure_formula: item for item in item_by_projected_name.values() if item.metric is not None } - out: List[OrderItem] = [] + out: list[OrderItem] = [] for ord_expr in order.args.get("expressions") or []: if not isinstance(ord_expr, exp.Ordered): continue @@ -1838,8 +1838,8 @@ def _translate_order_by( def _order_by_name( body: exp.Expression, *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> str: """Resolve an ORDER BY term to its projected name. @@ -1863,12 +1863,12 @@ def _order_by_name( def _validate_group_by( # NOSONAR(S3776) — single GROUP BY validation pass; extraction adds indirection - group: Optional[exp.Group], - derived: List[str], - item_by_projected_name: Dict[str, _ProjectionItem], + group: exp.Group | None, + derived: list[str], + item_by_projected_name: dict[str, _ProjectionItem], *, - strip_prefix: Optional[Tuple[str, str]] = None, - alias_map: Optional[Dict[str, str]] = None, + strip_prefix: tuple[str, str] | None = None, + alias_map: dict[str, str] | None = None, ) -> None: """Apply the strict-on-extras / lenient-on-omissions policy (§6.1). @@ -1881,7 +1881,7 @@ def _validate_group_by( # NOSONAR(S3776) — single GROUP BY validation pass; e if group is None: return derived_set = set(derived) - user_items: List[str] = [] + user_items: list[str] = [] for g in group.args.get("expressions") or []: if isinstance(g, exp.Column): user_items.append(_column_to_dotted(g, strip_prefix=strip_prefix, alias_map=alias_map)) @@ -1926,7 +1926,7 @@ def _is_ignorable_group_item(item: str) -> bool: def _reject_lossy_cast_in_implicit_grouping( - group: Optional[exp.Group], items: Sequence[_ProjectionItem], + group: exp.Group | None, items: Sequence[_ProjectionItem], ) -> None: """DEV-1566 — Codex round 12: when the user omits an explicit GROUP BY but projects at least one dimension, SLayer auto-groups (dim-only-dedup @@ -1969,7 +1969,7 @@ def _is_start_transaction(node: exp.Expression) -> bool: return body_name == "START" and alias_name == "TRANSACTION" -def _extract_set_setting(parsed: exp.Set) -> Optional[SetSettingOp]: +def _extract_set_setting(parsed: exp.Set) -> SetSettingOp | None: """Extract a ``SetSettingOp`` from a clean ``SET = `` / ``SET TO `` AST. Returns ``None`` for multi-item SETs or SetItem shapes whose body isn't a single ``EQ(Column(Identifier), rhs)`` @@ -2002,7 +2002,7 @@ def _extract_set_setting(parsed: exp.Set) -> Optional[SetSettingOp]: return SetSettingOp(name=name, value=value) -def _extract_setting_value(rhs: Optional[exp.Expression]) -> Optional[str]: +def _extract_setting_value(rhs: exp.Expression | None) -> str | None: """Return the string value of a ``SET`` rhs node, robust across the parser's literal vs. identifier vs. variable encodings: @@ -2039,7 +2039,7 @@ def _extract_setting_value(rhs: Optional[exp.Expression]) -> Optional[str]: return None -def _extract_reset_setting(parsed: exp.Command) -> Optional[ResetSettingOp]: +def _extract_reset_setting(parsed: exp.Command) -> ResetSettingOp | None: """Extract ``RESET `` or ``RESET ALL`` intent from an ``exp.Command`` root. Returns ``None`` for the bare ``RESET`` spelling — defensive fallback (drivers don't emit this in practice). @@ -2085,7 +2085,7 @@ def _command_expression_text(expr) -> str: return str(expr) -def _extract_command_form_set(parsed: exp.Command) -> Optional[SetSettingOp]: +def _extract_command_form_set(parsed: exp.Command) -> SetSettingOp | None: """For an ``exp.Command(this='SET', expression='')`` root, try to parse the raw expression text into a ``(name, value)`` pair. @@ -2124,7 +2124,7 @@ def _extract_command_form_set(parsed: exp.Command) -> Optional[SetSettingOp]: return SetSettingOp(name=name.lower(), value=value) -def _classify_noop_root(parsed: exp.Expression) -> Optional[NoOpResult]: +def _classify_noop_root(parsed: exp.Expression) -> NoOpResult | None: """Classify SET/SHOW/BEGIN/COMMIT/ROLLBACK roots into a NoOpResult with a facade-neutral ``command_tag``; ``None`` if not a no-op root. @@ -2173,7 +2173,7 @@ def _classify_noop_root(parsed: exp.Expression) -> Optional[NoOpResult]: def _unwrap_probe( - probe: "Union[RowBatch, ProbeMatcherOutcome]", + probe: "RowBatch | ProbeMatcherOutcome", ) -> ProbeResult: """Normalise a probe matcher's return value into a ``ProbeResult``. @@ -2194,10 +2194,10 @@ def translate( sql: str, catalog: FacadeCatalog, *, - dialect: Optional[str] = None, - probe_matcher: Optional[ProbeMatcher] = None, + dialect: str | None = None, + probe_matcher: ProbeMatcher | None = None, catalog_sql_executor: ( - "Optional[CatalogSqlExecutorProtocol | Callable[[], CatalogSqlExecutorProtocol]]" + "CatalogSqlExecutorProtocol | Callable[[], CatalogSqlExecutorProtocol] | None" ) = None, allow_column_cast: bool = True, ) -> TranslatorResult: @@ -2292,13 +2292,13 @@ class _ProjectionPlan(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - measures: List[dict] - dimension_refs: List[ColumnRef] - time_dims: List[TimeDimension] - time_dim_by_name: Dict[str, TimeDimension] - derived_dims: List[str] - column_name_mapping: List[Tuple[str, str]] - projection_types: List[Optional[DataType]] + measures: list[dict] + dimension_refs: list[ColumnRef] + time_dims: list[TimeDimension] + time_dim_by_name: dict[str, TimeDimension] + derived_dims: list[str] + column_name_mapping: list[tuple[str, str]] + projection_types: list[DataType | None] def _record_metric( @@ -2379,7 +2379,7 @@ def _build_projection_plan( def _index_items_by_canonical_form( items: Sequence[_ProjectionItem], -) -> Dict[str, _ProjectionItem]: +) -> dict[str, _ProjectionItem]: """Index projection items by alias plus canonical forms (time-grain only). Time-grain items are also keyed by ``grain(col)`` so an unaliased @@ -2397,7 +2397,7 @@ def _index_items_by_canonical_form( strict-on-extras error and the user aliases the CAST projection to reference it (``SELECT CAST(c AS T) AS x ... ORDER BY x``). """ - by_name: Dict[str, _ProjectionItem] = {item.projected_name: item for item in items} + by_name: dict[str, _ProjectionItem] = {item.projected_name: item for item in items} for item in items: if item.time_grain is not None and item.time_grain_underlying is not None: canonical = ( @@ -2407,7 +2407,7 @@ def _index_items_by_canonical_form( return by_name -def _parse_int_literal(node: Optional[exp.Expression]) -> Optional[int]: +def _parse_int_literal(node: exp.Expression | None) -> int | None: """Pull an int out of ``LIMIT N`` / ``OFFSET N`` style nodes.""" if node is None or not isinstance(node.expression, exp.Literal): return None @@ -2436,10 +2436,10 @@ class _JoinPlan(BaseModel): def _parse_left_join( *, - parsed_joins: List[exp.Join], + parsed_joins: list[exp.Join], parent_table: FacadeTable, catalog: FacadeCatalog, -) -> Optional[_JoinPlan]: +) -> _JoinPlan | None: """Validate and parse the single LEFT JOIN on the SELECT, if any. Returns ``None`` when no joins are present. Raises ``TranslationError`` @@ -2496,7 +2496,7 @@ def _reject_non_left_join(join: exp.Join) -> None: def _resolve_join_subquery_target( join: exp.Join, catalog: FacadeCatalog, -) -> Tuple[FacadeTable, str, str]: +) -> tuple[FacadeTable, str, str]: """Validate the right-side subquery shape and resolve its FROM table. Returns ``(target_facade_table, target_schema, join_alias)``. Raises @@ -2563,11 +2563,11 @@ def _resolve_join_subquery_target( def _parse_on_clause( *, - on: Optional[exp.Expression], + on: exp.Expression | None, parent_table: FacadeTable, target_table: FacadeTable, alias: str, -) -> Tuple[str, str]: +) -> tuple[str, str]: """Parse the ON clause as exactly one equality between a parent-table- qualified column and a join-alias-qualified column. @@ -2611,7 +2611,7 @@ def _parse_on_clause( ) -def _on_qualifier(col: exp.Column) -> Optional[str]: +def _on_qualifier(col: exp.Column) -> str | None: """The table-qualifier of an ON-clause column ref, or None if bare.""" table_node = col.args.get("table") if table_node is None: @@ -2620,7 +2620,7 @@ def _on_qualifier(col: exp.Column) -> Optional[str]: def _matches_parent_qualifier( - qual: Optional[str], parent_name: str, + qual: str | None, parent_name: str, ) -> bool: """True if the qualifier names the parent table (case-insensitive). Schema-prefix form is handled by ``_column_to_dotted``'s strip path, @@ -2628,7 +2628,7 @@ def _matches_parent_qualifier( return qual is not None and _ci_eq(qual, parent_name) -def _ci_eq(a: Optional[str], b: Optional[str]) -> bool: +def _ci_eq(a: str | None, b: str | None) -> bool: return a is not None and b is not None and a.lower() == b.lower() @@ -2691,7 +2691,7 @@ def _classify_against_parent_joins( target_name: str, source_col: str, target_col: str, -) -> Tuple[bool, bool]: +) -> tuple[bool, bool]: """Match the emitted (target_model, join_pairs) against the parent's configured joins. @@ -2727,9 +2727,9 @@ def _classify_against_parent_joins( def _materialise_dynamic_join_lookups( *, target_model: SlayerModel, - extra_dims_by_name: Dict[str, FacadeDimension], - extra_metrics_by_name: Dict[str, FacadeMetric], - extra_metrics_by_formula: Dict[str, FacadeMetric], + extra_dims_by_name: dict[str, FacadeDimension], + extra_metrics_by_name: dict[str, FacadeMetric], + extra_metrics_by_formula: dict[str, FacadeMetric], ) -> None: """For a dynamic-join target (no configured BFS expansion in the catalog), build the bare-col dims and col×agg metrics keyed by @@ -2808,14 +2808,14 @@ class _JoinOverlays(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - alias_map: Optional[Dict[str, str]] = None - extra_dims_by_name: Dict[str, FacadeDimension] = {} - extra_metrics_by_name: Dict[str, FacadeMetric] = {} - extra_metrics_by_formula: Dict[str, FacadeMetric] = {} + alias_map: dict[str, str] | None = None + extra_dims_by_name: dict[str, FacadeDimension] = {} + extra_metrics_by_name: dict[str, FacadeMetric] = {} + extra_metrics_by_formula: dict[str, FacadeMetric] = {} def _prepare_join_overlays( - join_plan: Optional[_JoinPlan], parent_name: str, + join_plan: _JoinPlan | None, parent_name: str, ) -> _JoinOverlays: """Build the alias map + (for dynamic joins) the materialised lookup overlays, and emit the join-related warnings. Extracted from @@ -2836,7 +2836,7 @@ def _prepare_join_overlays( return overlays -def _build_item_index(items: List[_ProjectionItem]) -> Dict[str, _ProjectionItem]: +def _build_item_index(items: list[_ProjectionItem]) -> dict[str, _ProjectionItem]: """Map every projection item by its user-facing name PLUS its canonical secondary key (time-grain canonical form for `CAST(date_trunc(...))` GROUP BY matches; dimension_ref dotted form for joined-col aliased @@ -2850,7 +2850,7 @@ def _build_item_index(items: List[_ProjectionItem]) -> Dict[str, _ProjectionItem explicitly named the bare column (and the engine ordering already matches that intent). """ - out: Dict[str, _ProjectionItem] = {item.projected_name: item for item in items} + out: dict[str, _ProjectionItem] = {item.projected_name: item for item in items} for item in items: if item.time_grain is not None and item.time_grain_underlying is not None: canonical = ( @@ -2886,7 +2886,7 @@ def _translate_slayer_select( # DEV-1558 B5: every helper that resolves a column ref needs the same # ``(schema, table)`` prefix-strip context as ``_resolve_projection``. - strip_prefix: Optional[Tuple[str, str]] = ( + strip_prefix: tuple[str, str] | None = ( (schema_name, table.name) if schema_name else None ) @@ -2916,7 +2916,7 @@ def _translate_slayer_select( ) _reject_lossy_cast_in_implicit_grouping(group, items) - filters: List[str] = [] + filters: list[str] = [] _apply_where( parsed.args.get("where"), plan.time_dim_by_name, item_by_projected_name, filters, diff --git a/slayer/flight/_capture_stub.py b/slayer/flight/_capture_stub.py index 27da282f..82d62173 100644 --- a/slayer/flight/_capture_stub.py +++ b/slayer/flight/_capture_stub.py @@ -21,7 +21,7 @@ import json import time from pathlib import Path -from typing import Any, Optional +from typing import Any import pyarrow as pa import pyarrow.flight as fl @@ -47,7 +47,7 @@ def _log(self, *, rpc: str, **payload: Any) -> None: f.write(json.dumps(record, default=str) + "\n") @staticmethod - def _b64(b: Optional[bytes]) -> Optional[str]: + def _b64(b: bytes | None) -> str | None: return base64.b64encode(b).decode("ascii") if b else None @staticmethod @@ -146,7 +146,7 @@ def list_actions(self, context: fl.ServerCallContext): return [] def do_action(self, context: fl.ServerCallContext, action: fl.Action): - body_bytes: Optional[bytes] = None + body_bytes: bytes | None = None if action.body is not None: body_bytes = action.body.to_pybytes() self._log( @@ -172,7 +172,7 @@ def _log(self, *, rpc: str, **payload: Any) -> None: f.write(json.dumps(record, default=str) + "\n") @staticmethod - def _b64(b: Optional[bytes]) -> Optional[str]: + def _b64(b: bytes | None) -> str | None: return base64.b64encode(b).decode("ascii") if b else None @staticmethod @@ -205,9 +205,9 @@ def __init__( location: str, handlers: FlightHandlers, log_path: Path, - token: Optional[str] = None, - tls_cert: Optional[str] = None, - tls_key: Optional[str] = None, + token: str | None = None, + tls_cert: str | None = None, + tls_key: str | None = None, ) -> None: super().__init__( location=location, @@ -244,7 +244,7 @@ def do_get(self, context: fl.ServerCallContext, ticket: fl.Ticket): return super().do_get(context, ticket) def do_action(self, context: fl.ServerCallContext, action: fl.Action): - body_bytes: Optional[bytes] = ( + body_bytes: bytes | None = ( action.body.to_pybytes() if action.body is not None else None ) self._recorder._log( diff --git a/slayer/flight/auth.py b/slayer/flight/auth.py index 6ed1d0e1..ce8c62bc 100644 --- a/slayer/flight/auth.py +++ b/slayer/flight/auth.py @@ -20,7 +20,6 @@ import hmac import ipaddress import logging -from typing import Optional import pyarrow.flight as fl @@ -52,7 +51,7 @@ def _is_loopback(host: str) -> bool: return False -def validate_bind_address(*, host: str, token: Optional[str]) -> None: +def validate_bind_address(*, host: str, token: str | None) -> None: """Raise ``ValueError`` if the server is about to bind a non-loopback address without a configured token (§4.3 / §7.1). """ @@ -66,7 +65,7 @@ def validate_bind_address(*, host: str, token: Optional[str]) -> None: ) -def validate_tls_pair(*, cert: Optional[str], key: Optional[str]) -> None: +def validate_tls_pair(*, cert: str | None, key: str | None) -> None: """TLS cert/key must be supplied together or not at all (§4.4).""" if (cert is None) != (key is None): raise ValueError( @@ -78,10 +77,10 @@ def validate_tls_pair(*, cert: Optional[str], key: Optional[str]) -> None: class _BearerTokenMiddleware(fl.ServerMiddleware): """No-op once-per-call middleware; auth check happened in the factory.""" - def __init__(self, *, environment_id: Optional[str] = None) -> None: + def __init__(self, *, environment_id: str | None = None) -> None: self._environment_id = environment_id - def call_completed(self, exception: Optional[BaseException]) -> None: + def call_completed(self, exception: BaseException | None) -> None: if exception is not None and self._environment_id is not None: logger.debug( "Flight SQL call (environmentId=%s) failed: %r", @@ -107,12 +106,12 @@ class BearerTokenMiddlewareFactory(fl.ServerMiddlewareFactory): handler-layer recheck would be possible if we ever want one.) """ - def __init__(self, *, token: Optional[str]) -> None: + def __init__(self, *, token: str | None) -> None: self._expected = token def start_call( self, info: fl.CallInfo, headers: dict - ) -> Optional[fl.ServerMiddleware]: + ) -> fl.ServerMiddleware | None: # Extract and lowercase header keys (gRPC standardises to lowercase # but client implementations differ). normalised = { @@ -121,7 +120,7 @@ def start_call( for k, v in (headers or {}).items() } env_id_raw = normalised.get("environmentid") - environment_id: Optional[str] = None + environment_id: str | None = None if isinstance(env_id_raw, (bytes, bytearray)): environment_id = env_id_raw.decode("utf-8", errors="replace") elif isinstance(env_id_raw, str): @@ -130,7 +129,7 @@ def start_call( logger.info("Flight SQL request environmentId=%s", environment_id) auth_raw = normalised.get("authorization") - provided: Optional[str] = None + provided: str | None = None if isinstance(auth_raw, (bytes, bytearray)): auth_raw = auth_raw.decode("utf-8", errors="replace") if isinstance(auth_raw, str) and auth_raw.lower().startswith("bearer "): diff --git a/slayer/flight/cli.py b/slayer/flight/cli.py index 14de3754..d472d008 100644 --- a/slayer/flight/cli.py +++ b/slayer/flight/cli.py @@ -10,7 +10,6 @@ import logging import os import sys -from typing import Optional logger = logging.getLogger(__name__) @@ -92,7 +91,7 @@ def run_flight_serve(args, *, resolve_storage, prepare_demo) -> None: engine = SlayerQueryEngine(storage=storage) handlers = FlightHandlers(engine=engine, storage=storage) - token: Optional[str] = args.token or os.environ.get("SLAYER_FLIGHT_TOKEN") + token: str | None = args.token or os.environ.get("SLAYER_FLIGHT_TOKEN") host = _resolve_host(host_arg=args.host, demo=args.demo, token=token) @@ -112,7 +111,7 @@ def run_flight_serve(args, *, resolve_storage, prepare_demo) -> None: server.serve() -def _resolve_host(*, host_arg: Optional[str], demo: bool, token: Optional[str]) -> str: +def _resolve_host(*, host_arg: str | None, demo: bool, token: str | None) -> str: """Apply the §7.1 demo-loopback default. If --host is not explicitly given AND --demo is set AND no token is diff --git a/slayer/flight/handlers.py b/slayer/flight/handlers.py index e8bf1fbc..aa9b514e 100644 --- a/slayer/flight/handlers.py +++ b/slayer/flight/handlers.py @@ -13,7 +13,6 @@ import decimal import logging from collections import defaultdict -from typing import Dict, List, Tuple import pyarrow as pa import pyarrow.flight as fl @@ -45,7 +44,7 @@ _TYPE_URL_PREFIX = "type.googleapis.com/arrow.flight.protocol.sql." -_COMMAND_BY_TYPE_URL: Dict[str, type] = { +_COMMAND_BY_TYPE_URL: dict[str, type] = { f"{_TYPE_URL_PREFIX}CommandStatementQuery": fsql_pb.CommandStatementQuery, f"{_TYPE_URL_PREFIX}CommandPreparedStatementQuery": fsql_pb.CommandPreparedStatementQuery, f"{_TYPE_URL_PREFIX}CommandGetCatalogs": fsql_pb.CommandGetCatalogs, @@ -67,7 +66,7 @@ } -def _decode_any(buf: bytes) -> Tuple[str, object]: +def _decode_any(buf: bytes) -> tuple[str, object]: """Decode an Any-wrapped Flight SQL command. Returns ``(type_url, message)``.""" any_msg = PbAny() any_msg.ParseFromString(buf) @@ -185,10 +184,10 @@ def _build_catalog(self) -> FlightCatalog: models_by_ds = self._fetch_models_by_datasource() return build_catalog(models_by_datasource=models_by_ds) - def _fetch_models_by_datasource(self) -> Dict[str, List[SlayerModel]]: - async def fetch() -> Dict[str, List[SlayerModel]]: + def _fetch_models_by_datasource(self) -> dict[str, list[SlayerModel]]: + async def fetch() -> dict[str, list[SlayerModel]]: datasources = await self._storage.list_datasources() - out: Dict[str, List[SlayerModel]] = defaultdict(list) + out: dict[str, list[SlayerModel]] = defaultdict(list) for ds in datasources: model_names = await self._storage.list_models(data_source=ds) for name in model_names: @@ -401,7 +400,7 @@ def _build_schema(result: "QueryResult") -> pa.Schema: @staticmethod def _rewrite_row( - row: dict, mapping: List[Tuple[str, str]], + row: dict, mapping: list[tuple[str, str]], ) -> dict: """Rewrite an engine row's keys into projected names + coerce Decimals.""" out: dict = {} @@ -439,12 +438,12 @@ def _build_flight_info( # --- top-level dispatch ------------------------------------------------------- -def decode_command(buf: bytes) -> Tuple[str, object]: +def decode_command(buf: bytes) -> tuple[str, object]: """Public re-export for tests / the server.""" return _decode_any(buf) -def decode_ticket(buf: bytes) -> Tuple[str, object]: +def decode_ticket(buf: bytes) -> tuple[str, object]: """Tickets are also Any-wrapped (TicketStatementQuery / CommandPreparedStatementQuery).""" return _decode_any(buf) diff --git a/slayer/flight/info_schema.py b/slayer/flight/info_schema.py index 3313e204..533d560f 100644 --- a/slayer/flight/info_schema.py +++ b/slayer/flight/info_schema.py @@ -7,7 +7,6 @@ from __future__ import annotations -from typing import Optional import pyarrow as pa import sqlglot.expressions as exp @@ -23,7 +22,7 @@ def match_info_schema( *, parsed: exp.Expression, catalog: FlightCatalog, -) -> Optional[pa.Table]: +) -> pa.Table | None: """Return the canned ``INFORMATION_SCHEMA.
`` answer as a ``pyarrow.Table`` or ``None``.""" batch = _shared_match_info_schema(parsed=parsed, catalog=catalog) diff --git a/slayer/flight/probe_queries.py b/slayer/flight/probe_queries.py index ebf22778..e963d038 100644 --- a/slayer/flight/probe_queries.py +++ b/slayer/flight/probe_queries.py @@ -6,7 +6,6 @@ from __future__ import annotations -from typing import Optional import pyarrow as pa import sqlglot.expressions as exp @@ -15,7 +14,7 @@ from slayer.flight.types import row_batch_to_arrow -def match_probe(parsed: exp.Expression) -> Optional[pa.Table]: +def match_probe(parsed: exp.Expression) -> pa.Table | None: """Return the canned ``pyarrow.Table`` for a matching probe, else ``None``.""" batch = _shared_match_probe(parsed) if batch is None: diff --git a/slayer/flight/server.py b/slayer/flight/server.py index b012e272..c3106bd1 100644 --- a/slayer/flight/server.py +++ b/slayer/flight/server.py @@ -11,7 +11,7 @@ import logging from pathlib import Path -from typing import Iterator, Optional +from collections.abc import Iterator import pyarrow.flight as fl @@ -89,9 +89,9 @@ def __init__( *, location: str, handlers: FlightHandlers, - token: Optional[str] = None, - tls_cert: Optional[str] = None, - tls_key: Optional[str] = None, + token: str | None = None, + tls_cert: str | None = None, + tls_key: str | None = None, ) -> None: tls_certificates = [] if tls_cert is not None and tls_key is not None: @@ -251,9 +251,9 @@ def build_server( host: str, port: int, handlers: FlightHandlers, - token: Optional[str] = None, - tls_cert: Optional[str] = None, - tls_key: Optional[str] = None, + token: str | None = None, + tls_cert: str | None = None, + tls_key: str | None = None, ) -> FlightSqlServer: """Factory wrapping ``FlightSqlServer`` with startup-time validation. diff --git a/slayer/flight/types.py b/slayer/flight/types.py index d3a2a926..453c7785 100644 --- a/slayer/flight/types.py +++ b/slayer/flight/types.py @@ -16,7 +16,6 @@ from __future__ import annotations -from typing import Optional import pyarrow as pa @@ -39,7 +38,7 @@ def datatype_to_arrow(dt: DataType) -> pa.DataType: return _DATATYPE_TO_ARROW[dt] -def arrow_to_datatype(at: pa.DataType) -> Optional[DataType]: +def arrow_to_datatype(at: pa.DataType) -> DataType | None: """Best-effort reverse map. Returns ``None`` if ``at`` cannot be coerced into one of the six diff --git a/slayer/help/__init__.py b/slayer/help/__init__.py index 68bb9d13..23185e8d 100644 --- a/slayer/help/__init__.py +++ b/slayer/help/__init__.py @@ -25,7 +25,6 @@ from __future__ import annotations from importlib.resources import files -from typing import Optional __all__ = ( "render_help", @@ -88,7 +87,7 @@ def available_topic_names() -> tuple[str, ...]: return _TOPIC_ORDER -def render_help(topic: Optional[str] = None) -> str: +def render_help(topic: str | None = None) -> str: """Render help content. * No topic (``None``, empty, or whitespace-only): return the intro body. diff --git a/slayer/mcp/server.py b/slayer/mcp/server.py index 22fe44d3..281173f8 100644 --- a/slayer/mcp/server.py +++ b/slayer/mcp/server.py @@ -2,7 +2,7 @@ import json import logging -from typing import Any, Dict, List, Optional, Tuple +from typing import Any import sqlalchemy as sa @@ -94,8 +94,8 @@ def _get_schemas(ds: DatasourceConfig) -> list[str]: def _fetch_tables( - ds: DatasourceConfig, schema_name: Optional[str] = None, -) -> Tuple[Optional[List[str]], Optional[str]]: + ds: DatasourceConfig, schema_name: str | None = None, +) -> tuple[list[str] | None, str | None]: """Inspect a datasource's table names. Returns ``(tables, None)`` on success or ``(None, friendly_error_message)`` @@ -164,7 +164,7 @@ def _cell_is_present(value: Any) -> bool: return True -def _truncate_description(text: Optional[str], max_chars: Optional[int]) -> Optional[str]: +def _truncate_description(text: str | None, max_chars: int | None) -> str | None: """Trim a description to ``max_chars`` and append the truncation marker. Returns the input unchanged when ``max_chars`` is ``None`` or the text is @@ -178,7 +178,7 @@ def _truncate_description(text: Optional[str], max_chars: Optional[int]) -> Opti return text[:max_chars] + _TRUNCATION_MARKER -def _format_meta(meta: Optional[Dict[str, Any]]) -> Optional[str]: +def _format_meta(meta: dict[str, Any] | None) -> str | None: """Compact JSON for the ``inspect_model`` meta cell. Returns ``None`` when ``meta`` is ``None`` so ``_markdown_table``'s @@ -190,8 +190,8 @@ def _format_meta(meta: Optional[Dict[str, Any]]) -> Optional[str]: def _resolve_inspect_sections( - sections: Optional[List[str]], -) -> Tuple[List[str], List[str]]: + sections: list[str] | None, +) -> tuple[list[str], list[str]]: """Validate and normalise the ``sections`` argument for ``inspect_model``. Returns ``(resolved, unknown)`` where ``resolved`` is the list of valid @@ -229,7 +229,7 @@ def _empty_ingest_message(*, schema_name: str, ds: DatasourceConfig) -> str: return "\n".join(lines) -def _render_new_models_section(new_models: List[Any]) -> List[str]: +def _render_new_models_section(new_models: list[Any]) -> list[str]: if not new_models: return [] lines = [f"Created {len(new_models)} new model(s):"] @@ -240,7 +240,7 @@ def _render_new_models_section(new_models: List[Any]) -> List[str]: return lines -def _render_updated_section(updated: List[Any]) -> List[str]: +def _render_updated_section(updated: list[Any]) -> list[str]: if not updated: return [] lines = [f"Updated {len(updated)} existing model(s):"] @@ -254,7 +254,7 @@ def _render_updated_section(updated: List[Any]) -> List[str]: return lines -def _render_unchanged_section(unchanged: List[Any]) -> List[str]: +def _render_unchanged_section(unchanged: list[Any]) -> list[str]: if not unchanged: return [] return [ @@ -263,7 +263,7 @@ def _render_unchanged_section(unchanged: List[Any]) -> List[str]: ] -def _render_drift_section(to_delete: List[Any]) -> List[str]: +def _render_drift_section(to_delete: list[Any]) -> list[str]: if not to_delete: return [] out = ["", "Pending drift (run validate_models / apply manually):"] @@ -271,7 +271,7 @@ def _render_drift_section(to_delete: List[Any]) -> List[str]: return out -def _render_errors_section(errors: List[Any]) -> List[str]: +def _render_errors_section(errors: list[Any]) -> list[str]: if not errors: return [] out = ["", f"Errors ({len(errors)}):"] @@ -307,7 +307,7 @@ def _render_ingest_result( if not a.created and not a.new_columns and not a.new_joins ] - lines: List[str] = [] + lines: list[str] = [] lines.extend(_render_new_models_section(new_models)) lines.extend(_render_updated_section(updated)) lines.extend(_render_unchanged_section(unchanged)) @@ -320,11 +320,11 @@ def _render_ingest_result( def _render_inspect_footer( *, - included: List[str], - names_only: List[str], - omitted: List[str], - unknown: List[str], -) -> Optional[str]: + included: list[str], + names_only: list[str], + omitted: list[str], + unknown: list[str], +) -> str | None: """Build the per-call truncation footer for ``inspect_model``. Returns ``None`` when there is nothing to report (no trimming, no @@ -332,7 +332,7 @@ def _render_inspect_footer( """ if not (names_only or omitted or unknown): return None - lines: List[str] = [] + lines: list[str] = [] if unknown: # repr() escapes newlines / quote chars so a caller-supplied value # like "foo\n> evil" can't forge additional footer lines. @@ -351,7 +351,7 @@ def _render_inspect_footer( return "\n".join(lines) if lines else None -def _markdown_table(rows: List[Dict[str, Any]], columns: List[str]) -> str: +def _markdown_table(rows: list[dict[str, Any]], columns: list[str]) -> str: """Render a list of row dicts as a GitHub-flavored markdown table. Columns with no present cell across every row are dropped automatically so @@ -393,8 +393,8 @@ def _markdown_table(rows: List[Dict[str, Any]], columns: List[str]) -> str: def _build_sample_query_args( model: SlayerModel, num_rows: int, - measure_types: Optional[Dict[str, str]] = None, -) -> Dict[str, Any]: + measure_types: dict[str, str] | None = None, +) -> dict[str, Any]: """Build the ``SlayerQuery`` payload for ``inspect_model``'s sample data. - First entry is always ``*:count``. @@ -412,7 +412,7 @@ def _build_sample_query_args( # Pick up to two categorical columns to group by first, so we don't also # aggregate them as measures (count_distinct(status) grouped by status is # always 1, which isn't useful sample data). - dims: List[Dict[str, str]] = [] + dims: list[dict[str, str]] = [] dim_names: set[str] = set() for c in model.columns: if c.hidden or c.primary_key: @@ -425,7 +425,7 @@ def _build_sample_query_args( if len(dims) >= 2: break - measures: List[Dict[str, str]] = [{"formula": "*:count"}] + measures: list[dict[str, str]] = [{"formula": "*:count"}] for c in model.columns: if c.hidden or c.primary_key or c.name in dim_names: continue @@ -460,10 +460,10 @@ def _build_sample_query_args( def _strip_model_prefix( - columns: List[str], - data: List[Dict[str, Any]], + columns: list[str], + data: list[dict[str, Any]], model_name: str, -) -> Tuple[List[str], List[Dict[str, Any]]]: +) -> tuple[list[str], list[dict[str, Any]]]: """Drop the redundant ``{model_name}.`` prefix from sample-data column keys. Keeps the markdown table compact (the model name already appears in the @@ -481,7 +481,7 @@ def _strip(key: str) -> str: async def _get_row_count( model: SlayerModel, engine: SlayerQueryEngine, -) -> Optional[int]: +) -> int | None: """Return the total row count of ``model``'s underlying table, or ``None`` on any failure. Uses a bare ``*:count`` query — the same aggregation a user would run to ask for the count. @@ -513,7 +513,7 @@ async def _get_row_count( async def _collect_measure_profile( model: SlayerModel, engine: SlayerQueryEngine, -) -> Dict[str, str]: +) -> dict[str, str]: """Probe min/max for each non-hidden, non-primary-key NUMERIC/TEMPORAL column via a single batched query. @@ -547,7 +547,7 @@ async def _collect_measure_profile( "type": str(c.type)} for c in columns ] - measures_payload: List[Dict[str, str]] = [] + measures_payload: list[dict[str, str]] = [] for c in columns: measures_payload.append({"formula": f"_slayer_probe_{c.name}:min"}) measures_payload.append({"formula": f"_slayer_probe_{c.name}:max"}) @@ -562,7 +562,7 @@ async def _collect_measure_profile( except Exception: return {} - result: Dict[str, str] = {} + result: dict[str, str] = {} for c in columns: mn = row.get(f"{model.name}._slayer_probe_{c.name}_min") mx = row.get(f"{model.name}._slayer_probe_{c.name}_max") @@ -573,7 +573,7 @@ async def _collect_measure_profile( return result -def _build_backing_query_info(model: SlayerModel) -> Optional[dict]: +def _build_backing_query_info(model: SlayerModel) -> dict | None: """Build the ``backing_query`` block for inspect_model output. Returns ``None`` for non-query-backed models. For query-backed models, @@ -588,7 +588,7 @@ def _build_backing_query_info(model: SlayerModel) -> Optional[dict]: from slayer.core.query import extract_placeholder_names all_placeholders: set = set() - stage_dicts: List[dict] = [] + stage_dicts: list[dict] = [] # A placeholder is "required" only if it has no default at any layer the # engine consults: model.query_variables OR the stage's own variables. defaulted: set = set(model.query_variables.keys()) @@ -635,7 +635,7 @@ def _render_stage_field_list(key: str, val: list) -> str: return "; ".join(_render_field_value(v) for v in val) -def _render_source_model(src: Any) -> Optional[str]: +def _render_source_model(src: Any) -> str | None: """Render a stage's ``source_model`` (str or ModelExtension dict).""" if isinstance(src, str): return f"- source_model: `{src}`" @@ -646,10 +646,10 @@ def _render_source_model(src: Any) -> Optional[str]: return None -def _render_stage(i: int, stage: dict, total: int) -> List[str]: +def _render_stage(i: int, stage: dict, total: int) -> list[str]: """Render one stage's markdown lines.""" title = stage.get("name") or ("final" if i == total else f"stage {i}") - out: List[str] = [f"\n**{i}. {title}**"] + out: list[str] = [f"\n**{i}. {title}**"] src_line = _render_source_model(stage.get("source_model")) if src_line: out.append(src_line) @@ -663,7 +663,7 @@ def _render_stage(i: int, stage: dict, total: int) -> List[str]: def _backing_query_markdown_section(info: dict) -> str: """Format the ``backing_query`` info as a markdown section.""" - lines: List[str] = ["## Backing Query"] + lines: list[str] = ["## Backing Query"] stages = info.get("stages") or [] for i, stage in enumerate(stages, start=1): lines.extend(_render_stage(i, stage, len(stages))) @@ -766,25 +766,25 @@ def create_mcp_server( # NOSONAR(S3776) — FastMCP tool-registration factory; ) @mcp.tool(description=_help_description) - async def help(topic: Optional[str] = None) -> str: # noqa: A001 — intentional shadow of builtin inside factory + async def help(topic: str | None = None) -> str: # noqa: A001 — intentional shadow of builtin inside factory return render_help(topic=topic) @mcp.tool() async def query( # NOSONAR S107 — FastMCP introspects this signature to expose each query option as a typed MCP tool argument; collapsing into a dict would degrade the agent-facing schema source_model: str | ModelExtension | SlayerModel, - measures: Optional[List[Dict[str, str]]] = None, - dimensions: Optional[List[str]] = None, - filters: Optional[List[str]] = None, - time_dimensions: Optional[List[Dict[str, Any]]] = None, - order: Optional[List[Dict[str, str]]] = None, - limit: Optional[int] = None, - offset: Optional[int] = None, + measures: list[dict[str, str]] | None = None, + dimensions: list[str] | None = None, + filters: list[str] | None = None, + time_dimensions: list[dict[str, Any]] | None = None, + order: list[dict[str, str]] | None = None, + limit: int | None = None, + offset: int | None = None, whole_periods_only: bool = False, show_sql: bool = False, dry_run: bool = False, explain: bool = False, format: str = "markdown", - variables: Optional[Dict[str, Any]] = None, + variables: dict[str, Any] | None = None, distinct_dimension_values: bool = True, ) -> str: """Query data from a semantic model. Call inspect_model first to see available columns and measures. @@ -826,7 +826,7 @@ async def query( # NOSONAR S107 — FastMCP introspects this signature to expos Before calling this tool, run ``search`` first, supplying the entities you're thinking of using (and/or the query itself via the ``query`` arg, or a free-text ``question``). Read the returned memories and consider any matching example queries before formulating the final query. """ - data: Dict[str, Any] = {"source_model": source_model} + data: dict[str, Any] = {"source_model": source_model} if dimensions: data["dimensions"] = list(dimensions) if filters: @@ -917,8 +917,8 @@ async def query( # NOSONAR S107 — FastMCP introspects this signature to expos @mcp.tool() async def query_nested( - queries: List[Dict[str, Any]], - variables: Optional[Dict[str, Any]] = None, + queries: list[dict[str, Any]], + variables: dict[str, Any] | None = None, show_sql: bool = False, dry_run: bool = False, explain: bool = False, @@ -1039,7 +1039,7 @@ async def models_summary( return f"Datasource '{datasource_name}' not found." all_names = await storage.list_models(data_source=datasource_name) - matched: List[SlayerModel] = [] + matched: list[SlayerModel] = [] for n in all_names: try: m = await storage.get_model(n, data_source=datasource_name) @@ -1100,11 +1100,11 @@ async def models_summary( indent=2, ) - sections: List[str] = [ + sections: list[str] = [ f"# Datasource: `{datasource_name}` — {len(matched)} model(s)" ] for m in matched: - model_lines: List[str] = [f"## `{m.name}`"] + model_lines: list[str] = [f"## `{m.name}`"] if m.description: model_lines.append(m.description) @@ -1163,9 +1163,9 @@ async def inspect_model( num_rows: int = 3, show_sql: bool = False, format: str = "markdown", - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, - data_source: Optional[str] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, + data_source: str | None = None, compact: bool = True, ) -> str: """Return a complete-yet-compact view of a semantic model. @@ -1257,12 +1257,12 @@ async def inspect_model( ] truncated_model_desc = _truncate_description(model.description, descriptions_max_chars) - out_sections: List[str] = [f"# Model: `{model.name}`"] + out_sections: list[str] = [f"# Model: `{model.name}`"] if truncated_model_desc: out_sections.append(truncated_model_desc) # Metadata bullets (incl. row_count from a cheap *:count query) - meta: List[str] = [] + meta: list[str] = [] if model.data_source: meta.append(f"- **data_source:** `{model.data_source}`") if model.sql_table: @@ -1307,12 +1307,12 @@ async def inspect_model( # / ``distinct_count`` columns of the row. Read the cache first; on # miss, profile live and write back via storage so subsequent calls # (and any search) hit the cache for free (DEV-1375 + DEV-1480). - profile_by_name: Dict[str, str] = {} - profile_values_by_name: Dict[str, Optional[List[str]]] = {} - distinct_count_by_name: Dict[str, Optional[int]] = {} - measure_profile: Dict[str, str] = {} + profile_by_name: dict[str, str] = {} + profile_values_by_name: dict[str, list[str] | None] = {} + distinct_count_by_name: dict[str, int | None] = {} + measure_profile: dict[str, str] = {} if "columns" in included_set: - uncached_columns: List[Column] = [] + uncached_columns: list[Column] = [] for c in model.columns: if c.hidden or c.primary_key: continue @@ -1354,9 +1354,9 @@ async def inspect_model( async def _persist_sample( *, col_name: str, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: try: await storage.update_column_sampled( @@ -1454,7 +1454,7 @@ async def _persist_sample( # ``measure_types`` informs the sample query's choice of avg vs # count_distinct. Only needed when ``samples`` is in the included set. - measure_types: Dict[str, str] = {} + measure_types: dict[str, str] = {} if "samples" in included_set: measure_types = await engine.get_column_types( model_name=model.name, @@ -1466,7 +1466,7 @@ async def _persist_sample( # ------------------------------------------------------------------ visible_columns = [c for c in model.columns if not c.hidden] if "columns" in included_set: - col_rows: List[Dict[str, Any]] = [] + col_rows: list[dict[str, Any]] = [] for c in visible_columns: aggs = ", ".join(c.allowed_aggregations) if c.allowed_aggregations else "all" # DEV-1480: key-presence check (not ``or`` truthiness) so an @@ -1509,7 +1509,7 @@ async def _persist_sample( # Measures section # ------------------------------------------------------------------ if "measures" in included_set: - measure_rows: List[Dict[str, Any]] = [] + measure_rows: list[dict[str, Any]] = [] for mm in model.measures: measure_rows.append({ "name": mm.name, @@ -1536,7 +1536,7 @@ async def _persist_sample( # ------------------------------------------------------------------ if "aggregations" in included_set: if model.aggregations: - agg_rows: List[Dict[str, Any]] = [] + agg_rows: list[dict[str, Any]] = [] for a in model.aggregations: if a.params: if show_sql: @@ -1571,7 +1571,7 @@ async def _persist_sample( # Joins section # ------------------------------------------------------------------ if "joins" in included_set: - join_rows: List[Dict[str, Any]] = [] + join_rows: list[dict[str, Any]] = [] for j in model.joins: pairs = "; ".join(f"{src} = {tgt}" for src, tgt in j.join_pairs) join_rows.append({ @@ -1594,9 +1594,9 @@ async def _persist_sample( # ------------------------------------------------------------------ # Sample data (fully omitted when not in sections) # ------------------------------------------------------------------ - sample_sql: Optional[str] = None - sample_data: Optional[Dict[str, Any]] = None - sample_error: Optional[str] = None + sample_sql: str | None = None + sample_data: dict[str, Any] | None = None + sample_error: str | None = None if "samples" in included_set: query_args = _build_sample_query_args( model=model, num_rows=num_rows, measure_types=measure_types, @@ -1641,8 +1641,8 @@ async def _persist_sample( # ``None``; query-bearing memories are recall-only. Auto-pruned when # no learning-shaped memory matches. # ------------------------------------------------------------------ - relevant_learnings: List[Any] = [] - wanted: List[str] = [] + relevant_learnings: list[Any] = [] + wanted: list[str] = [] if "learnings" in included_set: ds = model.data_source wanted = [f"{ds}.{model.name}"] @@ -1690,7 +1690,7 @@ async def _persist_sample( ) if fmt == "json": - payload: Dict[str, Any] = { + payload: dict[str, Any] = { "model_name": model.name, "description": truncated_model_desc, "data_source": model.data_source, @@ -1712,7 +1712,7 @@ async def _persist_sample( # Columns if "columns" in included_set: - col_payloads: List[Dict[str, Any]] = [] + col_payloads: list[dict[str, Any]] = [] for c in visible_columns: # DEV-1480 key-presence (not ``or`` truthiness) so empty # string ``sampled=""`` (all-NULL categorical) survives. @@ -1858,14 +1858,14 @@ async def _persist_sample( @mcp.tool() async def create_model( name: str, - sql_table: Optional[str] = None, - sql: Optional[str] = None, - data_source: Optional[str] = None, - description: Optional[str] = None, - columns: Optional[List[Dict[str, Any]]] = None, - measures: Optional[List[Dict[str, Any]]] = None, - query: Optional[Any] = None, - variables: Optional[Dict[str, Any]] = None, + sql_table: str | None = None, + sql: str | None = None, + data_source: str | None = None, + description: str | None = None, + columns: list[dict[str, Any]] | None = None, + measures: list[dict[str, Any]] | None = None, + query: Any | None = None, + variables: dict[str, Any] | None = None, ) -> str: """Create a new semantic model, either from a database table or from a query. @@ -1964,7 +1964,7 @@ def _upsert_entity( id_field: str, changes: list, label: str, - ) -> Optional[str]: + ) -> str | None: """Upsert a named entity in *entity_list*. Returns an error string on validation failure, ``None`` on success. @@ -1999,23 +1999,23 @@ def _upsert_entity( @mcp.tool() async def edit_model( model_name: str, - description: Optional[str] = None, - data_source: Optional[str] = None, - new_data_source: Optional[str] = None, - default_time_dimension: Optional[str] = None, - sql_table: Optional[str] = None, - sql: Optional[str] = None, - source_queries: Optional[List[Dict[str, Any]]] = None, + description: str | None = None, + data_source: str | None = None, + new_data_source: str | None = None, + default_time_dimension: str | None = None, + sql_table: str | None = None, + sql: str | None = None, + source_queries: list[dict[str, Any]] | None = None, query_variables: Any = _UNSET, - hidden: Optional[bool] = None, - columns: Optional[List[Dict[str, Any]]] = None, - measures: Optional[List[Dict[str, Any]]] = None, - aggregations: Optional[List[Dict[str, Any]]] = None, - joins: Optional[List[Dict[str, Any]]] = None, - add_filters: Optional[List[str]] = None, - remove_filters: Optional[List[str]] = None, - remove: Optional[Dict[str, List[str]]] = None, - meta: Optional[Dict[str, Any]] = _UNSET, + hidden: bool | None = None, + columns: list[dict[str, Any]] | None = None, + measures: list[dict[str, Any]] | None = None, + aggregations: list[dict[str, Any]] | None = None, + joins: list[dict[str, Any]] | None = None, + add_filters: list[str] | None = None, + remove_filters: list[str] | None = None, + remove: dict[str, list[str]] | None = None, + meta: dict[str, Any] | None = _UNSET, ) -> str: """Edit an existing model in a single call — update metadata, upsert columns/measures/aggregations/joins, manage filters, and remove entities. @@ -2082,7 +2082,7 @@ async def edit_model( return f"Model '{model_name}' not found." original_data_source = model.data_source - changes: List[str] = [] + changes: list[str] = [] # DEV-1375: track refresh-triggering changes so the post-save hook # knows whether to refresh just the touched columns or every # column on the model. @@ -2338,7 +2338,7 @@ async def edit_model( # changed the indexed text. Best-effort: any raise here is # captured into ``refresh_warnings`` so the save's success # status survives a flaky embedding API. - refresh_warnings: List[str] = [] + refresh_warnings: list[str] = [] if changed_columns or model_level_change or model_doc_changed: try: refresh_warnings = await handle_edit_refresh( @@ -2375,13 +2375,13 @@ async def edit_model( async def create_datasource( name: str, type: str, - host: Optional[str] = None, - port: Optional[int] = None, - database: Optional[str] = None, - username: Optional[str] = None, - password: Optional[str] = None, - connection_string: Optional[str] = None, - schema_name: Optional[str] = None, + host: str | None = None, + port: int | None = None, + database: str | None = None, + username: str | None = None, + password: str | None = None, + connection_string: str | None = None, + schema_name: str | None = None, auto_ingest: bool = True, ) -> str: """Create a database connection, verify it, and auto-ingest models. Use ${ENV_VAR} syntax in credentials to reference environment variables. @@ -2541,7 +2541,7 @@ async def describe_datasource( @mcp.tool() async def edit_datasource( name: str, - description: Optional[str] = None, + description: str | None = None, ) -> str: """Update a datasource's metadata. @@ -2569,9 +2569,9 @@ async def edit_datasource( # round-7 review: the refresh is post-save and best-effort — # log a warning if it raises and surface a partial-success # message rather than telling the agent the save itself failed. - refresh_warning: Optional[str] = None + refresh_warning: str | None = None if description is not None and description != old_description: - models_in_ds: List[SlayerModel] = [] + models_in_ds: list[SlayerModel] = [] for model_name in await storage.list_models(data_source=name): m = await storage.get_model(model_name, data_source=name) if m is not None: @@ -2599,7 +2599,7 @@ async def edit_datasource( # ----------------------------------------------------------------------- @mcp.tool() - async def delete_model(name: str, data_source: Optional[str] = None) -> str: + async def delete_model(name: str, data_source: str | None = None) -> str: """Delete a semantic model. Args: @@ -2617,7 +2617,7 @@ async def delete_model(name: str, data_source: Optional[str] = None) -> str: return f"Model '{name}' not found." @mcp.tool() - async def validate_models(data_source: Optional[str] = None) -> str: + async def validate_models(data_source: str | None = None) -> str: """Diff persisted SLayer models against the live database schema(s). Returns a JSON-serialized list of pending delete operations @@ -2697,7 +2697,7 @@ async def ingest_datasource_models(datasource_name: str, include_tables: str = " ) @mcp.tool() - async def set_datasource_priority(priority: List[str]) -> str: + async def set_datasource_priority(priority: list[str]) -> str: """Configure how SLayer disambiguates bare model names that exist in multiple datasources. @@ -2743,8 +2743,8 @@ def _format_resolution_error(exc: Exception) -> str: async def save_memory( learning: str, linked_entities: Any, - id: Optional[str] = None, # noqa: A002 — MCP arg name - description: Optional[str] = None, + id: str | None = None, # noqa: A002 — MCP arg name + description: str | None = None, ) -> str: """Save an agent memory: a free-form note plus the SLayer entities it concerns. @@ -2853,12 +2853,12 @@ async def forget_memory(id: Any) -> str: # noqa: A002 — MCP arg name @mcp.tool() async def search( - entities: Optional[List[str]] = None, + entities: list[str] | None = None, query: Any = None, - question: Optional[str] = None, - datasource: Optional[str] = None, + question: str | None = None, + datasource: str | None = None, max_results: int = 10, - cypher_filter: Optional[str] = None, + cypher_filter: str | None = None, compact: bool = True, ) -> str: """Up to three-channel semantic search over memories + canonical entities. @@ -2931,12 +2931,12 @@ async def search( return mcp -def _build_dict(**kwargs: Any) -> Dict[str, Any]: +def _build_dict(**kwargs: Any) -> dict[str, Any]: """Build a dict from keyword arguments, excluding None values.""" return {k: v for k, v in kwargs.items() if v is not None} -def _format_table(data: List[Dict[str, Any]], columns: List[str], max_rows: int = 50) -> str: +def _format_table(data: list[dict[str, Any]], columns: list[str], max_rows: int = 50) -> str: """Format data as a pipe-separated table (used for sample data display).""" if not data: return "No results." @@ -2956,14 +2956,14 @@ def _format_table(data: List[Dict[str, Any]], columns: List[str], max_rows: int return result -def _format_json(data: List[Dict[str, Any]], columns: List[str]) -> str: +def _format_json(data: list[dict[str, Any]], columns: list[str]) -> str: """Format data as JSON array.""" import json return json.dumps(data, default=str) -def _format_csv(data: List[Dict[str, Any]], columns: List[str]) -> str: +def _format_csv(data: list[dict[str, Any]], columns: list[str]) -> str: """Format data as CSV.""" if not data: return "" @@ -2988,7 +2988,7 @@ def _format_output(result: SlayerResponse, fmt: str) -> str: return _format_json(data=result.data, columns=result.columns) -def _format_field_meta(entries: Dict[str, Any]) -> List[str]: +def _format_field_meta(entries: dict[str, Any]) -> list[str]: """Format a dict of field metadata entries into lines.""" lines = [] for col, fm in entries.items(): diff --git a/slayer/memories/models.py b/slayer/memories/models.py index d4024845..1ca6df7b 100644 --- a/slayer/memories/models.py +++ b/slayer/memories/models.py @@ -19,7 +19,7 @@ """ from datetime import datetime, timezone -from typing import Any, List, Optional +from typing import Any from pydantic import BaseModel, Field, field_validator, model_validator @@ -87,9 +87,9 @@ class Memory(BaseModel): version: int = 2 id: str = "" learning: str - description: Optional[str] = None - entities: List[str] = Field(default_factory=list) - query: Optional[SlayerQuery] = None + description: str | None = None + entities: list[str] = Field(default_factory=list) + query: SlayerQuery | None = None created_at: datetime = Field(default_factory=_utcnow) @model_validator(mode="before") @@ -143,7 +143,7 @@ def _normalise_description(cls, value: Any) -> Any: @field_validator("description") @classmethod - def _check_description_length(cls, value: Optional[str]) -> Optional[str]: + def _check_description_length(cls, value: str | None) -> str | None: """DEV-1549: hard cap so a single memory hit can never balloon the search payload.""" if value is not None and len(value) > MEMORY_DESCRIPTION_MAX_CHARS: @@ -161,8 +161,8 @@ def _check_description_length(cls, value: Optional[str]) -> Optional[str]: class SaveMemoryResponse(BaseModel): memory_id: str - resolved_entities: List[str] - warnings: List[str] = Field(default_factory=list) + resolved_entities: list[str] + warnings: list[str] = Field(default_factory=list) class ForgetMemoryResponse(BaseModel): diff --git a/slayer/memories/ranker.py b/slayer/memories/ranker.py index 5702abfd..f7d46f8a 100644 --- a/slayer/memories/ranker.py +++ b/slayer/memories/ranker.py @@ -29,7 +29,6 @@ from __future__ import annotations -from typing import List, Tuple from rank_bm25 import BM25Plus @@ -37,9 +36,9 @@ def bm25_rank( - memories: List[Memory], - query_entities: List[str], -) -> List[Tuple[Memory, float]]: + memories: list[Memory], + query_entities: list[str], +) -> list[tuple[Memory, float]]: """Rank ``memories`` against ``query_entities`` using BM25Plus. Returns ``(memory, score)`` pairs sorted by score descending. @@ -68,7 +67,7 @@ def bm25_rank( bm25 = BM25Plus(tokenised) scores = bm25.get_scores(list(query_set)) - paired: List[Tuple[Memory, float]] = [ + paired: list[tuple[Memory, float]] = [ (memories[i], float(scores[i])) for i in eligible ] paired.sort(key=lambda pair: pair[1], reverse=True) diff --git a/slayer/memories/resolver.py b/slayer/memories/resolver.py index 3d676516..4be127d6 100644 --- a/slayer/memories/resolver.py +++ b/slayer/memories/resolver.py @@ -29,7 +29,7 @@ from __future__ import annotations import re -from typing import Iterable, List, Optional, Set, Tuple +from collections.abc import Iterable from pydantic import BaseModel @@ -104,8 +104,8 @@ class EntityResolution(BaseModel): datasource-vs-model collision; fatal failures raise instead. """ - canonical_forms: List[str] - warnings: List[str] = [] + canonical_forms: list[str] + warnings: list[str] = [] def _model_has_leaf(model: SlayerModel, leaf: str) -> bool: @@ -122,9 +122,9 @@ def _model_has_leaf(model: SlayerModel, leaf: str) -> bool: async def _all_models_in_datasource( storage: StorageBackend, data_source: str -) -> List[SlayerModel]: +) -> list[SlayerModel]: identities = await storage._list_all_model_identities() - out: List[SlayerModel] = [] + out: list[SlayerModel] = [] for ds, name in identities: if ds != data_source: continue @@ -136,7 +136,7 @@ async def _all_models_in_datasource( async def _find_leaf_in_priority_winner( storage: StorageBackend, leaf: str -) -> Tuple[Optional[str], List[SlayerModel]]: +) -> tuple[str | None, list[SlayerModel]]: """Walk the priority list; return ``(data_source, matches)`` for the first datasource that has ≥1 model carrying ``leaf`` as a column / measure / custom aggregation. ``matches`` may have multiple models @@ -156,7 +156,7 @@ async def _find_leaf_in_priority_winner( async def _resolve_join_path( storage: StorageBackend, starting_model: SlayerModel, - path: List[str], + path: list[str], ) -> SlayerModel: """Walk a chain of join targets, returning the leaf model. @@ -188,7 +188,7 @@ async def _resolve_join_path( async def _resolve_dotted_against_model( storage: StorageBackend, starting_model: SlayerModel, - rest: List[str], + rest: list[str], ) -> str: """Apply the leaf rule to a path ``[hop, hop, ..., leaf?]`` rooted at ``starting_model``. Returns the canonical form.""" @@ -223,7 +223,7 @@ async def resolve_entity( # NOSONAR(S3776) — single linear dispatch matching raw: str, *, storage: StorageBackend, - source_model: Optional[SlayerModel] = None, + source_model: SlayerModel | None = None, ) -> EntityResolution: """Resolve a single entity reference. @@ -300,7 +300,7 @@ async def resolve_entity( # NOSONAR(S3776) — single linear dispatch matching f"'{raw}' contains an invalid identifier segment." ) - warnings: List[str] = [] + warnings: list[str] = [] known_dses = set(await storage.list_datasources()) # ----- step 3: datasource-prefix detection --------------------------- @@ -431,7 +431,7 @@ def _formula_aggregated_refs(field: FieldSpec) -> Iterable[AggregatedMeasureRef] _FILTER_TOKEN_RE = re.compile(r"[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*)*") -def _extract_filter_tokens(filter_text: str) -> List[str]: +def _extract_filter_tokens(filter_text: str) -> list[str]: """Extract identifier-shaped tokens from a filter expression that might be entity references. @@ -442,7 +442,7 @@ def _extract_filter_tokens(filter_text: str) -> List[str]: cleaned = _FILTER_AGG_SUFFIX_RE.sub("", filter_text) cleaned = _FILTER_LITERAL_RE.sub("", cleaned) cleaned = _FILTER_VAR_RE.sub("", cleaned) - out: List[str] = [] + out: list[str] = [] for m in _FILTER_TOKEN_RE.finditer(cleaned): token = m.group(0) # Skip identifiers immediately followed by '(' — they're SQL @@ -474,9 +474,9 @@ async def extract_entities_from_query( # NOSONAR(S3776) — straight-line walk always tagged, even if no field references it explicitly. Resolution failures bubble up unchanged. """ - canonical: List[str] = [] - warnings: List[str] = [] - seen: Set[str] = set() + canonical: list[str] = [] + warnings: list[str] = [] + seen: set[str] = set() def _add(forms: Iterable[str]) -> None: for f in forms: @@ -584,7 +584,7 @@ def _add(forms: Iterable[str]) -> None: warnings.extend(result.warnings) # Deduplicate warnings while preserving order. - seen_warn: Set[str] = set() + seen_warn: set[str] = set() deduped_warnings = [] for w in warnings: if w not in seen_warn: diff --git a/slayer/memories/service.py b/slayer/memories/service.py index 1fa4e4c5..f856ad09 100644 --- a/slayer/memories/service.py +++ b/slayer/memories/service.py @@ -20,7 +20,6 @@ from __future__ import annotations -from typing import List, Optional, Union from slayer.core.query import SlayerQuery from slayer.memories.models import ( @@ -35,8 +34,8 @@ from slayer.storage.base import StorageBackend -QueryInput = Union[SlayerQuery, dict] -LinkedEntities = Union[List[str], SlayerQuery, dict] +QueryInput = SlayerQuery | dict +LinkedEntities = list[str] | SlayerQuery | dict def _coerce_query(query: QueryInput) -> SlayerQuery: @@ -57,7 +56,7 @@ def _coerce_query(query: QueryInput) -> SlayerQuery: ) -def _coerce_memory_id(identifier: Union[int, str]) -> str: +def _coerce_memory_id(identifier: int | str) -> str: """DEV-1428: accept native ``str`` (canonical form) or legacy ``int`` (back-compat: stringify decimally). Validates the result through :func:`_validate_memory_id_charset` so the surface layer @@ -79,9 +78,9 @@ def _coerce_memory_id(identifier: Union[int, str]) -> str: return value -def _dedup(items: List[str]) -> List[str]: +def _dedup(items: list[str]) -> list[str]: seen: set[str] = set() - out: List[str] = [] + out: list[str] = [] for x in items: if x not in seen: seen.add(x) @@ -104,17 +103,17 @@ async def save_memory( *, learning: str, linked_entities: LinkedEntities, - id: Optional[str] = None, # noqa: A002 — public kwarg - description: Optional[str] = None, + id: str | None = None, # noqa: A002 — public kwarg + description: str | None = None, ) -> SaveMemoryResponse: if not learning or not learning.strip(): raise ValueError("learning text must be a non-empty string.") if id is not None: _validate_memory_id_charset(id) - canonical: List[str] = [] - warnings: List[str] = [] - attached_query: Optional[SlayerQuery] = None + canonical: list[str] = [] + warnings: list[str] = [] + attached_query: SlayerQuery | None = None if isinstance(linked_entities, list): if not linked_entities: @@ -169,7 +168,7 @@ async def save_memory( # ---- forget_memory ------------------------------------------------- async def forget_memory( - self, *, identifier: Union[int, str] + self, *, identifier: int | str ) -> ForgetMemoryResponse: memory_id = _coerce_memory_id(identifier) await self._storage.delete_memory(memory_id) diff --git a/slayer/pg_facade/auth.py b/slayer/pg_facade/auth.py index 33666d81..ce04c3e4 100644 --- a/slayer/pg_facade/auth.py +++ b/slayer/pg_facade/auth.py @@ -16,7 +16,6 @@ import hmac import ipaddress import logging -from typing import Optional logger = logging.getLogger(__name__) @@ -41,7 +40,7 @@ def _is_loopback(host: str) -> bool: return any(ip in net for net in _LOOPBACK_NETWORKS) -def validate_bind_address(*, host: str, token: Optional[str]) -> None: +def validate_bind_address(*, host: str, token: str | None) -> None: """Raise ``ValueError`` if binding a non-loopback address without a token.""" if token: return @@ -53,7 +52,7 @@ def validate_bind_address(*, host: str, token: Optional[str]) -> None: ) -def validate_tls_pair(*, cert: Optional[str], key: Optional[str]) -> None: +def validate_tls_pair(*, cert: str | None, key: str | None) -> None: """TLS cert/key must be supplied together or not at all.""" if (cert is None) != (key is None): raise ValueError( @@ -62,7 +61,7 @@ def validate_tls_pair(*, cert: Optional[str], key: Optional[str]) -> None: ) -def verify_password(client_password: str, expected: Optional[str]) -> bool: +def verify_password(client_password: str, expected: str | None) -> bool: """Constant-time cleartext-password check. When no token is configured (``expected is None``) any non-empty password diff --git a/slayer/pg_facade/connection.py b/slayer/pg_facade/connection.py index 55a1ac4e..0ee05741 100644 --- a/slayer/pg_facade/connection.py +++ b/slayer/pg_facade/connection.py @@ -18,7 +18,7 @@ import logging import re import struct -from typing import Dict, Iterable, List, Optional, Set, Tuple +from collections.abc import Iterable import sqlglot import sqlglot.errors @@ -78,7 +78,7 @@ # ``integer_datetimes``, ``is_superuser``, ``in_hot_standby``, # ``default_transaction_read_only`` are GUC_REPORT in real Postgres too but # the facade doesn't expose them as settable, so we don't list them here. -_GUC_REPORT_NAMES: Dict[str, str] = { +_GUC_REPORT_NAMES: dict[str, str] = { "application_name": "application_name", "client_encoding": "client_encoding", "datestyle": "DateStyle", @@ -93,18 +93,18 @@ # DEV-1570 type aliases — populated by ``_build_column_type_index`` below and # cached per-connection. Declared here so the ``__init__`` annotation can # reference them without a forward-ref dance. -ColumnTypeKey = Tuple[str, str, str] # (schema_lower, table_lower, column_lower) -ColumnTypeIndex = Dict[ColumnTypeKey, DataType] +ColumnTypeKey = tuple[str, str, str] # (schema_lower, table_lower, column_lower) +ColumnTypeIndex = dict[ColumnTypeKey, DataType] class _PreparedStatement(BaseModel): sql: str - parameter_oids: List[int] + parameter_oids: list[int] class _Portal(BaseModel): sql: str - result_format_codes: List[int] + result_format_codes: list[int] class _Done(Exception): @@ -121,7 +121,7 @@ def __init__( *, engine, storage, - token: Optional[str], + token: str | None, tls_ctx=None, ) -> None: self._reader = reader @@ -131,15 +131,15 @@ def __init__( self._token = token self._tls_ctx = tls_ctx self._tx_state: bytes = proto.TX_IDLE - self._datasource: Optional[str] = None - self._catalog: Optional[FacadeCatalog] = None - self._statements: Dict[str, _PreparedStatement] = {} - self._portals: Dict[str, _Portal] = {} + self._datasource: str | None = None + self._catalog: FacadeCatalog | None = None + self._statements: dict[str, _PreparedStatement] = {} + self._portals: dict[str, _Portal] = {} # Lazily-built (schema, table, column) -> DataType lookup, used by the # DEV-1570 empty-string-vs-non-text Bind rewrite. Built once per # connection on first need; ``None`` until then so connections that # never bind candidates pay zero cost. - self._column_type_index: Optional[ColumnTypeIndex] = None + self._column_type_index: ColumnTypeIndex | None = None # Extended protocol: after an error the backend discards every message # until the next Sync, then resumes with ReadyForQuery. self._skip_until_sync = False @@ -147,7 +147,7 @@ def __init__( # set_config writes; consulted by SHOW / current_setting reads. Seeded # from the module-level SESSION_SETTING_SEED via dict(...) so each # connection owns its own copy (never aliasing the seed). - self._session_settings: Dict[str, str] = dict(SESSION_SETTING_SEED) + self._session_settings: dict[str, str] = dict(SESSION_SETTING_SEED) # DEV-1569: when True, ``_describe_sql`` is in flight — translator # calls must remain pure. Suppresses application of any session- # setting mutation hints surfaced by the probe matcher during a @@ -175,7 +175,7 @@ async def run(self) -> None: # ----- startup ---------------------------------------------------------- - async def _read_startup_frame(self) -> Optional[bytes]: + async def _read_startup_frame(self) -> bytes | None: """Read a startup-style frame (no type byte). Returns the body (starting with the 4-byte code) or ``None`` on EOF / malformed length.""" try: @@ -191,7 +191,7 @@ async def _read_startup_frame(self) -> Optional[bytes]: except asyncio.IncompleteReadError: return None - async def _handle_startup(self) -> Optional[proto.StartupMessage]: + async def _handle_startup(self) -> proto.StartupMessage | None: while True: body = await self._read_startup_frame() if body is None: @@ -267,7 +267,7 @@ async def _authenticate(self) -> bool: # ----- datasource resolution ------------------------------------------- - async def _resolve_datasource(self, database: Optional[str]) -> bool: + async def _resolve_datasource(self, database: str | None) -> bool: datasources = await self._storage.list_datasources() if database and database in datasources: self._datasource = database @@ -282,7 +282,7 @@ async def _resolve_datasource(self, database: Optional[str]) -> bool: async def _build_catalog(self) -> FacadeCatalog: assert self._datasource is not None - models: List[SlayerModel] = [] + models: list[SlayerModel] = [] names = await self._storage.list_models(data_source=self._datasource) for name in names: model = await self._storage.get_model(name=name, data_source=self._datasource) @@ -455,7 +455,7 @@ def _substitute_params(self, stmt: _PreparedStatement, bind: proto.BindMessage) empty_string_null_params = self._empty_string_null_params_for_bind( sql=stmt.sql, raw_values=bind.parameter_values, oids=oids, ) - literals: List[str] = [] + literals: list[str] = [] for i, (raw, fmt, oid) in enumerate( zip(bind.parameter_values, formats, oids), start=1, ): @@ -480,8 +480,8 @@ def repl(match: "re.Match[str]") -> str: return _PARAM_PLACEHOLDER.sub(repl, stmt.sql) def _empty_string_null_params_for_bind( - self, *, sql: str, raw_values, oids: List[int], - ) -> Set[int]: + self, *, sql: str, raw_values, oids: list[int], + ) -> set[int]: # DEV-1570: pre-classify $N indices whose bound value is an empty # text-OID payload AND whose AST occurrence targets a non-TEXT catalog # column. Those positions emit ``NULL`` rather than ``''`` so DuckDB @@ -533,8 +533,8 @@ async def _handle_describe(self, msg: proto.DescribeMessage) -> None: ) def _describe_sql( - self, sql: str, *, result_formats: Optional[List[int]], - param_oids: Optional[List[int]] = None, + self, sql: str, *, result_formats: list[int] | None, + param_oids: list[int] | None = None, ) -> None: # DEV-1558 fix: the catalog executor's Describe path runs the SQL # against DuckDB to obtain the cursor's column description. When the @@ -653,7 +653,7 @@ def _probe_matcher(self, parsed: exp.Expression): return facade_match_probe(parsed) async def _run_statement( - self, sql: str, *, result_formats: Optional[List[int]], send_row_description: bool, + self, sql: str, *, result_formats: list[int] | None, send_row_description: bool, ) -> bool: """Translate + respond. Returns False if an error was sent.""" try: @@ -696,7 +696,7 @@ async def _run_statement( return False def _emit_row_batch( - self, batch: RowBatch, result_formats: Optional[List[int]], send_row_description: bool, + self, batch: RowBatch, result_formats: list[int] | None, send_row_description: bool, ) -> None: formats = proto.parse_result_format_codes(result_formats or [], len(batch.columns)) if send_row_description: @@ -724,7 +724,7 @@ def _emit_row_batch( self._writer.write(proto.encode_command_complete(f"SELECT {len(batch.rows)}")) async def _run_query( - self, result: QueryResult, result_formats: Optional[List[int]], send_row_description: bool, + self, result: QueryResult, result_formats: list[int] | None, send_row_description: bool, ) -> bool: try: response = await self._engine.execute( @@ -757,8 +757,8 @@ async def _run_query( return True def _fields_for_result( - self, result, result_formats: Optional[List[int]], - ) -> Optional[List[proto.FieldDescription]]: + self, result, result_formats: list[int] | None, + ) -> list[proto.FieldDescription] | None: if isinstance(result, (ProbeResult, InfoSchemaResult, PgCatalogResult)): cols = result.batch.columns formats = proto.parse_result_format_codes(result_formats or [], len(cols)) @@ -782,7 +782,7 @@ def _fields_for_result( # ----- transaction state ------------------------------------------------- - def _apply_tx_command(self, command_tag: Optional[str]) -> None: + def _apply_tx_command(self, command_tag: str | None) -> None: if command_tag in ("BEGIN", "START TRANSACTION"): self._tx_state = proto.TX_IN_TRANSACTION elif command_tag in ("COMMIT", "ROLLBACK", "END"): @@ -882,7 +882,7 @@ async def _flush(self) -> None: # declaring text OIDs for parameters that compared against int columns # (``objsubid = $N``), which the literal ``''`` sentinel turned into # an unanswerable text-vs-int comparison. -_TYPED_SENTINEL_BY_OID: Dict[int, str] = { +_TYPED_SENTINEL_BY_OID: dict[int, str] = { proto.OID_TEXT: "CAST(NULL AS VARCHAR)", proto.OID_INT8: "CAST(NULL AS BIGINT)", proto.OID_FLOAT8: "CAST(NULL AS DOUBLE)", @@ -892,7 +892,7 @@ async def _flush(self) -> None: } -def _substitute_typed_sentinels(sql: str, param_oids: List[int]) -> str: +def _substitute_typed_sentinels(sql: str, param_oids: list[int]) -> str: """Replace each ``$N`` placeholder with a typed sentinel literal derived from ``param_oids[N-1]``. Falls back to bare ``NULL`` when the OID is unknown (e.g. extra placeholders past the declared list) @@ -1019,7 +1019,7 @@ def _classify_empty_string_param_targets( *, sql: str, column_type_index: ColumnTypeIndex, candidate_param_indices: Iterable[int], -) -> Set[int]: +) -> set[int]: """Return the subset of ``candidate_param_indices`` whose AST occurrences appear in a comparison / IN / BETWEEN predicate against a column that resolves via ``column_type_index`` to a non-TEXT ``DataType``. @@ -1063,9 +1063,9 @@ def _classify_empty_string_param_targets( def _map_column_tables( *, parsed: exp.Expression, column_type_index: ColumnTypeIndex, -) -> Dict[int, Optional[Tuple[str, str]]]: +) -> dict[int, tuple[str, str] | None]: """Resolve every Column node to its owning scope's (schema, table).""" - column_to_table: Dict[int, Optional[Tuple[str, str]]] = {} + column_to_table: dict[int, tuple[str, str] | None] = {} for scope in traverse_scope(parsed): sources = _resolved_table_sources(scope.sources) for col in scope.find_all(exp.Column): @@ -1078,13 +1078,13 @@ def _map_column_tables( def _collect_non_text_params( *, parsed: exp.Expression, - candidates: Set[int], - column_to_table: Dict[int, Optional[Tuple[str, str]]], + candidates: set[int], + column_to_table: dict[int, tuple[str, str] | None], column_type_index: ColumnTypeIndex, -) -> Set[int]: +) -> set[int]: """Walk comparison / IN / BETWEEN nodes; collect $N indices whose paired Column resolves to a non-TEXT ``DataType``.""" - result: Set[int] = set() + result: set[int] = set() for node in parsed.walk(): for col, param_idx in _column_param_pairs_from_node(node): if param_idx not in candidates: @@ -1099,7 +1099,7 @@ def _collect_non_text_params( return result -def _column_param_pairs_from_node(node) -> Iterable[Tuple[exp.Column, int]]: +def _column_param_pairs_from_node(node) -> Iterable[tuple[exp.Column, int]]: """Yield (Column, param_index) pairs for each comparison-shape node.""" if isinstance(node, _COMPARISON_NODE_TYPES): yield from _pair_column_and_param(node.this, node.expression) @@ -1117,7 +1117,7 @@ def _column_param_pairs_from_node(node) -> Iterable[Tuple[exp.Column, int]]: yield from _pair_column_and_param(value, el) -def _try_extract_column(node) -> Optional[exp.Column]: +def _try_extract_column(node) -> exp.Column | None: """Return the underlying ``exp.Column`` if ``node`` is one (optionally wrapped in semantically-transparent ``exp.Paren`` layers). CAST / function / arithmetic wrappers around a column are documented out of scope — @@ -1129,7 +1129,7 @@ def _try_extract_column(node) -> Optional[exp.Column]: return node if isinstance(node, exp.Column) else None -def _try_extract_param_index(node) -> Optional[int]: +def _try_extract_param_index(node) -> int | None: """Return the 1-based ``$N`` index if ``node`` is an ``exp.Parameter`` (optionally wrapped in ``exp.Paren`` and/or ``exp.Cast`` layers). @@ -1151,7 +1151,7 @@ def _try_extract_param_index(node) -> Optional[int]: return None -def _pair_column_and_param(left, right) -> Iterable[Tuple[exp.Column, int]]: +def _pair_column_and_param(left, right) -> Iterable[tuple[exp.Column, int]]: """Two operands where one is a bare Column and the other is a $N Parameter -> yield (Column, $N). Returns nothing if both sides are the same kind or if neither is a Column / Parameter.""" @@ -1169,7 +1169,7 @@ def _pair_column_and_param(left, right) -> Iterable[Tuple[exp.Column, int]]: yield col, param_idx -def _resolved_table_sources(sources) -> Dict[str, Tuple[str, str]]: +def _resolved_table_sources(sources) -> dict[str, tuple[str, str]]: """Given ``Scope.sources``, return ``{alias_lower: (schema_lower, table_lower)}``. Sources whose value is another ``Scope`` (CTE / derived subquery) @@ -1181,13 +1181,13 @@ def _resolved_table_sources(sources) -> Dict[str, Tuple[str, str]]: information_schema requires an explicit qualifier — bare names like ``columns`` never resolve there (Codex round 1, finding #4). """ - result: Dict[str, Tuple[str, str]] = {} + result: dict[str, tuple[str, str]] = {} for alias, src in sources.items(): if not isinstance(src, exp.Table): continue tbl_name = src.name.lower() db_part = src.args.get("db") - schema: Optional[str] = None + schema: str | None = None if db_part is not None: schema_raw = db_part.name if hasattr(db_part, "name") else str(db_part) schema = schema_raw.lower() @@ -1201,9 +1201,9 @@ def _resolved_table_sources(sources) -> Dict[str, Tuple[str, str]]: def _resolve_column_table( *, col: exp.Column, - scope_sources: Dict[str, Tuple[str, str]], + scope_sources: dict[str, tuple[str, str]], column_type_index: ColumnTypeIndex, -) -> Optional[Tuple[str, str]]: +) -> tuple[str, str] | None: """Resolve a Column node to ``(schema, table_name)`` via its owning scope's sources. Returns ``None`` if unresolvable (no scope match, ambiguous bare name, or table not in scope).""" @@ -1216,7 +1216,7 @@ def _resolve_column_table( return scope_sources[table_q] return None name_lower = col.name.lower() - matches: List[Tuple[str, str]] = [] + matches: list[tuple[str, str]] = [] for _alias, (schema, tbl) in scope_sources.items(): if (schema, tbl, name_lower) in column_type_index: matches.append((schema, tbl)) @@ -1225,7 +1225,7 @@ def _resolve_column_table( return None -def _resolve_param_oids(stmt: _PreparedStatement) -> List[int]: +def _resolve_param_oids(stmt: _PreparedStatement) -> list[int]: """The parameter OIDs to report in ParameterDescription. asyncpg leaves ``Parse`` parameter OIDs empty and relies on the server to @@ -1253,7 +1253,7 @@ def _is_tx_end(stmt: exp.Expression) -> bool: return False -def _command_tag(command_tag: Optional[str]) -> str: +def _command_tag(command_tag: str | None) -> str: if command_tag in ("BEGIN", "START TRANSACTION"): return "BEGIN" if command_tag is None: @@ -1272,7 +1272,7 @@ def _sqlstate_for(exc: TranslationError) -> str: return proto.SQLSTATE_FEATURE_NOT_SUPPORTED -def _encode_value(value, oid: int, fmt: int) -> Optional[bytes]: +def _encode_value(value, oid: int, fmt: int) -> bytes | None: if fmt == proto.FORMAT_BINARY: return value_to_binary(value, oid) return value_to_text(value, oid) diff --git a/slayer/pg_facade/identity.py b/slayer/pg_facade/identity.py index 051b38dc..ea54bc2d 100644 --- a/slayer/pg_facade/identity.py +++ b/slayer/pg_facade/identity.py @@ -8,7 +8,6 @@ from __future__ import annotations -from typing import List, Tuple import slayer @@ -23,7 +22,7 @@ def version_string() -> str: ) -def parameter_status_defaults() -> List[Tuple[str, str]]: +def parameter_status_defaults() -> list[tuple[str, str]]: """The ParameterStatus burst sent after auth, before the first ReadyForQuery. UTC + UTF8 keep timestamp / encoding handling unambiguous.""" return [ diff --git a/slayer/pg_facade/probes.py b/slayer/pg_facade/probes.py index 8aa7e20a..de8e4dce 100644 --- a/slayer/pg_facade/probes.py +++ b/slayer/pg_facade/probes.py @@ -24,7 +24,6 @@ from __future__ import annotations -from typing import Dict, Optional import sqlglot.expressions as exp @@ -38,7 +37,7 @@ # are case-insensitive). Values align with what the startup # ``ParameterStatus`` burst (identity.py) advertises for the same setting, # so a client SHOW immediately after connect agrees with the burst. -SESSION_SETTING_SEED: Dict[str, str] = { +SESSION_SETTING_SEED: dict[str, str] = { "search_path": '"$user", public', "transaction_isolation": "read committed", "standard_conforming_strings": "on", @@ -63,14 +62,14 @@ } -def _single(name: str, value: Optional[str], dtype: DataType = DataType.TEXT) -> RowBatch: +def _single(name: str, value: str | None, dtype: DataType = DataType.TEXT) -> RowBatch: return RowBatch( columns=[FacadeColumn(name=name, type=dtype)], rows=[{name: value}], ) -def _single_projection(parsed: exp.Expression) -> Optional[exp.Expression]: +def _single_projection(parsed: exp.Expression) -> exp.Expression | None: if not isinstance(parsed, exp.Select): return None exprs = parsed.args.get("expressions") or [] @@ -82,7 +81,7 @@ def _single_projection(parsed: exp.Expression) -> Optional[exp.Expression]: return body -def _show_setting_name(parsed: exp.Expression) -> Optional[str]: +def _show_setting_name(parsed: exp.Expression) -> str | None: if not isinstance(parsed, exp.Command): return None if str(parsed.this).upper() != "SHOW": @@ -94,7 +93,7 @@ def _show_setting_name(parsed: exp.Expression) -> Optional[str]: return name.strip().strip("'\"") -def _anonymous_name(node: exp.Expression) -> Optional[str]: +def _anonymous_name(node: exp.Expression) -> str | None: if isinstance(node, exp.Anonymous): return str(node.this).lower() return None @@ -102,8 +101,8 @@ def _anonymous_name(node: exp.Expression) -> Optional[str]: def match_pg_probe( parsed: exp.Expression, *, datasource: str, version_str: str, - session_settings: Optional[Dict[str, str]] = None, -) -> Optional[RowBatch]: + session_settings: dict[str, str] | None = None, +) -> RowBatch | None: """Return a datasource-aware canned ``RowBatch`` for a Postgres probe, else ``None`` (caller falls back to the shared probe matcher). @@ -156,8 +155,8 @@ def match_pg_probe( def match_pg_probe_with_mutation( parsed: exp.Expression, *, datasource: str, version_str: str, - session_settings: Optional[Dict[str, str]] = None, -) -> Optional[ProbeMatcherOutcome]: + session_settings: dict[str, str] | None = None, +) -> ProbeMatcherOutcome | None: """Mutation-aware variant of :func:`match_pg_probe`. Returns a :class:`ProbeMatcherOutcome` carrying both the row batch and (for ``set_config(name, value, ...)`` matches) a :class:`SetSettingOp` @@ -180,7 +179,7 @@ def match_pg_probe_with_mutation( return ProbeMatcherOutcome(batch=batch, settings_mutation=mutation) -def _extract_set_config_mutation(parsed: exp.Expression) -> Optional[SetSettingOp]: +def _extract_set_config_mutation(parsed: exp.Expression) -> SetSettingOp | None: """Inspect a parsed AST root for ``SELECT set_config('', '', )`` and return a ``SetSettingOp`` carrying the (lowercased name, raw value) pair; return ``None`` otherwise. @@ -253,7 +252,7 @@ def _is_truthy_boolean_literal(value: str) -> bool: ) -def _first_literal(node: exp.Anonymous, index: int) -> Optional[str]: +def _first_literal(node: exp.Anonymous, index: int) -> str | None: """Return the string value of the ``index``-th argument of ``node`` if it is a string literal (or a CAST around a string literal). Returns ``None`` otherwise. @@ -274,7 +273,7 @@ def _first_literal(node: exp.Anonymous, index: int) -> Optional[str]: return None -def _setting_value(node: exp.Anonymous, settings: Dict[str, str]) -> str: +def _setting_value(node: exp.Anonymous, settings: dict[str, str]) -> str: """``current_setting('')`` → ``settings[name]`` (lowercased lookup); unknown settings return the empty string.""" setting = (_first_literal(node, 0) or "").lower() diff --git a/slayer/pg_facade/protocol.py b/slayer/pg_facade/protocol.py index 68cd091b..309639a1 100644 --- a/slayer/pg_facade/protocol.py +++ b/slayer/pg_facade/protocol.py @@ -16,7 +16,6 @@ from __future__ import annotations import struct -from typing import Dict, List, Optional, Tuple from pydantic import BaseModel @@ -110,7 +109,7 @@ class FieldDescription(BaseModel): column_attr: int = 0 -def encode_row_description(fields: List[FieldDescription]) -> bytes: +def encode_row_description(fields: list[FieldDescription]) -> bytes: payload = struct.pack(">h", len(fields)) for f in fields: payload += _cstr(f.name) @@ -126,7 +125,7 @@ def encode_row_description(fields: List[FieldDescription]) -> bytes: return _msg(b"T", payload) -def encode_data_row(values: List[Optional[bytes]]) -> bytes: +def encode_data_row(values: list[bytes | None]) -> bytes: payload = struct.pack(">h", len(values)) for v in values: if v is None: @@ -164,7 +163,7 @@ def encode_portal_suspended() -> bytes: return _msg(b"s", b"") -def encode_parameter_description(oids: List[int]) -> bytes: +def encode_parameter_description(oids: list[int]) -> bytes: payload = struct.pack(">h", len(oids)) for oid in oids: payload += struct.pack(">i", oid) @@ -195,21 +194,21 @@ def encode_notice_response(*, code: str, message: str, severity: str = "NOTICE") class StartupMessage(BaseModel): protocol_version: int - parameters: Dict[str, str] + parameters: dict[str, str] class ParseMessage(BaseModel): name: str query: str - parameter_oids: List[int] + parameter_oids: list[int] class BindMessage(BaseModel): portal: str statement: str - parameter_format_codes: List[int] - parameter_values: List[Optional[bytes]] - result_format_codes: List[int] + parameter_format_codes: list[int] + parameter_values: list[bytes | None] + result_format_codes: list[int] class DescribeMessage(BaseModel): @@ -275,7 +274,7 @@ def decode_startup(body: bytes) -> StartupMessage: null terminator.""" r = _Reader(buf=body) version = r.int32() - params: Dict[str, str] = {} + params: dict[str, str] = {} while r.pos < len(body): if body[r.pos:r.pos + 1] == b"\x00": break @@ -311,7 +310,7 @@ def decode_bind(body: bytes) -> BindMessage: n_fmt = _nonneg_count(r.int16(), "parameter format code") fmt_codes = [r.int16() for _ in range(n_fmt)] n_params = _nonneg_count(r.int16(), "parameter") - values: List[Optional[bytes]] = [] + values: list[bytes | None] = [] for _ in range(n_params): length = r.int32() if length == -1: @@ -355,13 +354,13 @@ def decode_close(body: bytes) -> CloseMessage: # --- generic message splitting (used by the connection layer + tests) -------- -def split_messages(buf: bytes) -> List[Tuple[str, bytes]]: +def split_messages(buf: bytes) -> list[tuple[str, bytes]]: """Split a byte stream of tagged messages into ``[(type_char, body), …]``. Used by tests to verify encoded server messages and by any client-side helper. The body excludes the type byte and the length prefix. """ - out: List[Tuple[str, bytes]] = [] + out: list[tuple[str, bytes]] = [] pos = 0 while pos < len(buf): type_char = buf[pos:pos + 1].decode("ascii") @@ -372,14 +371,14 @@ def split_messages(buf: bytes) -> List[Tuple[str, bytes]]: return out -def validate_format_codes(codes: List[int]) -> None: +def validate_format_codes(codes: list[int]) -> None: """Reject format codes outside ``{text, binary}`` (protocol violation).""" for c in codes: if c not in (FORMAT_TEXT, FORMAT_BINARY): raise ValueError(f"invalid format code {c!r} (must be 0=text or 1=binary)") -def parse_result_format_codes(codes: List[int], column_count: int) -> List[int]: +def parse_result_format_codes(codes: list[int], column_count: int) -> list[int]: """Resolve Bind result-format codes to one entry per result column. Per the protocol: 0 codes → all text; 1 code → applies to every column; diff --git a/slayer/pg_facade/server.py b/slayer/pg_facade/server.py index d667e378..9b22fac5 100644 --- a/slayer/pg_facade/server.py +++ b/slayer/pg_facade/server.py @@ -14,7 +14,6 @@ import os import ssl import sys -from typing import Optional from slayer.pg_facade.auth import validate_bind_address, validate_tls_pair from slayer.pg_facade.connection import PgConnection @@ -30,8 +29,8 @@ async def serve( port: int, engine, storage, - token: Optional[str] = None, - tls_ctx: Optional[ssl.SSLContext] = None, + token: str | None = None, + tls_ctx: ssl.SSLContext | None = None, ) -> None: """Bind and serve forever. Validates the bind/token combination first.""" validate_bind_address(host=host, token=token) @@ -54,7 +53,7 @@ async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> await server.serve_forever() -def _build_tls_context(cert: Optional[str], key: Optional[str]) -> Optional[ssl.SSLContext]: +def _build_tls_context(cert: str | None, key: str | None) -> ssl.SSLContext | None: validate_tls_pair(cert=cert, key=key) if cert is None or key is None: return None @@ -64,12 +63,12 @@ def _build_tls_context(cert: Optional[str], key: Optional[str]) -> Optional[ssl. return ctx -def _resolve_token(token_arg: Optional[str]) -> Optional[str]: +def _resolve_token(token_arg: str | None) -> str | None: """``--token`` wins over the ``$SLAYER_PG_TOKEN`` env var.""" return token_arg or os.environ.get("SLAYER_PG_TOKEN") -def _resolve_host(*, host_arg: Optional[str], demo: bool, token: Optional[str]) -> str: +def _resolve_host(*, host_arg: str | None, demo: bool, token: str | None) -> str: """Explicit --host wins; --demo without a token defaults to loopback so the no-token fallback applies; otherwise bind all interfaces.""" if host_arg is not None: @@ -92,7 +91,7 @@ def run_pg_serve(args, *, resolve_storage, prepare_demo) -> None: prepare_demo(args, storage) engine = SlayerQueryEngine(storage=storage) - token: Optional[str] = _resolve_token(args.token) + token: str | None = _resolve_token(args.token) host = _resolve_host(host_arg=args.host, demo=args.demo, token=token) try: diff --git a/slayer/pg_facade/types.py b/slayer/pg_facade/types.py index 301b38f3..162c068e 100644 --- a/slayer/pg_facade/types.py +++ b/slayer/pg_facade/types.py @@ -23,7 +23,7 @@ import math import struct from decimal import Decimal -from typing import Any, Dict, Optional +from typing import Any from slayer.core.enums import DataType from slayer.pg_facade.protocol import ( @@ -35,7 +35,7 @@ OID_TIMESTAMP, ) -DATATYPE_TO_OID: Dict[DataType, int] = { +DATATYPE_TO_OID: dict[DataType, int] = { DataType.TEXT: OID_TEXT, DataType.INT: OID_INT8, DataType.DOUBLE: OID_FLOAT8, @@ -49,7 +49,7 @@ _PG_EPOCH_DATETIME = _dt.datetime(2000, 1, 1) -def datatype_to_oid(dt: Optional[DataType]) -> int: +def datatype_to_oid(dt: DataType | None) -> int: """Map a SLayer ``DataType`` to a Postgres OID; unknown / None → text.""" if dt is None: return OID_TEXT @@ -61,7 +61,7 @@ def datatype_to_oid(dt: Optional[DataType]) -> int: def value_to_text( # NOSONAR(S3776) — flat per-Python-type dispatch value: Any, oid: int = OID_TEXT, -) -> Optional[bytes]: +) -> bytes | None: """Encode an engine value as Postgres text-format bytes (``None`` → SQL NULL). ``oid`` lets the encoder coerce values to match the declared column type. @@ -124,7 +124,7 @@ def _format_timestamp(value: _dt.datetime) -> str: # --- binary-format output ---------------------------------------------------- -def value_to_binary(value: Any, oid: int) -> Optional[bytes]: +def value_to_binary(value: Any, oid: int) -> bytes | None: """Encode an engine value as Postgres binary-format bytes for ``oid`` (``None`` → SQL NULL).""" if value is None: diff --git a/slayer/search/cypher_naive.py b/slayer/search/cypher_naive.py index 0dada6fc..94f2cb47 100644 --- a/slayer/search/cypher_naive.py +++ b/slayer/search/cypher_naive.py @@ -11,7 +11,6 @@ from __future__ import annotations import re -from typing import Set from slayer.core.errors import SlayerError @@ -38,7 +37,7 @@ _AS_ID_RE = re.compile(pattern=r"\bAS\s+id\b", flags=re.IGNORECASE) -def parse_naive_label_filter(cypher: str) -> Set[str]: +def parse_naive_label_filter(cypher: str) -> set[str]: """Parse a simple MATCH (n:Label1:Label2) RETURN n.id AS id expression and return the set of kind strings to filter search results on. @@ -62,7 +61,7 @@ def parse_naive_label_filter(cypher: str) -> Set[str]: ) labels_str = match.group("labels") labels = [lb.strip() for lb in re.split(r"\s*:\s*", labels_str) if lb.strip()] - kinds: Set[str] = set() + kinds: set[str] = set() for label in labels: kind = _LABEL_TO_KIND.get(label.lower()) if kind is None: diff --git a/slayer/search/graph.py b/slayer/search/graph.py index 4db0055e..61880c55 100644 --- a/slayer/search/graph.py +++ b/slayer/search/graph.py @@ -39,7 +39,7 @@ import os import re from functools import lru_cache -from typing import Any, Dict, FrozenSet, Optional +from typing import Any from pydantic import BaseModel @@ -227,7 +227,7 @@ def _insert_nodes( datasource_names: list[str], visible_models: dict, memories: list, - datasource_descriptions: Optional[Dict[str, Optional[str]]] = None, + datasource_descriptions: dict[str, str | None] | None = None, ) -> None: """Insert all node rows into the graph. @@ -445,7 +445,7 @@ async def build_graph(storage: StorageBackend) -> tuple[Any, Any]: # DEV-1549: load datasource descriptions so the graph schema can # expose `Datasource.description` for cypher_filter queries. - datasource_descriptions: Dict[str, Optional[str]] = {} + datasource_descriptions: dict[str, str | None] = {} for ds_name in datasource_names: cfg = await storage.get_datasource(ds_name) datasource_descriptions[ds_name] = ( @@ -531,7 +531,7 @@ async def _get_or_rebuild(storage: StorageBackend) -> tuple[Any, Any]: key = _storage_key(storage) try: - current_fp: Optional[str] = await storage.graph_fingerprint() + current_fp: str | None = await storage.graph_fingerprint() except OSError: current_fp = None @@ -562,7 +562,7 @@ async def _get_or_rebuild(storage: StorageBackend) -> tuple[Any, Any]: async def get_filtered_ids( cypher: str, storage: StorageBackend, -) -> FrozenSet[str]: +) -> frozenset[str]: """Execute a Cypher query against the storage graph and return the frozenset of id strings from the result's ``id`` column. diff --git a/slayer/search/index.py b/slayer/search/index.py index ee7a13c1..ddd58485 100644 --- a/slayer/search/index.py +++ b/slayer/search/index.py @@ -19,7 +19,6 @@ from __future__ import annotations -from typing import Dict, List, Optional, Tuple import tantivy from pydantic import BaseModel, ConfigDict @@ -48,7 +47,7 @@ class IndexHit(BaseModel): canonical: str text: str score: float - memory_id: Optional[str] = None # populated only when kind == "memory" + memory_id: str | None = None # populated only when kind == "memory" # --------------------------------------------------------------------------- @@ -83,10 +82,10 @@ def _add_doc( def build_in_memory_index( *, - memories: List[Memory], - models: List[SlayerModel], - datasources: List[str], - datasource_descriptions: Optional[Dict[str, Optional[str]]] = None, + memories: list[Memory], + models: list[SlayerModel], + datasources: list[str], + datasource_descriptions: dict[str, str | None] | None = None, ) -> tantivy.Index: """Build a fresh in-RAM tantivy index covering the corpus. @@ -124,18 +123,18 @@ class Corpus(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) index: "tantivy.Index" - canonical_to_text: Dict[str, str] - canonical_to_kind: Dict[str, str] - canonical_to_description: Dict[str, Optional[str]] = {} + canonical_to_text: dict[str, str] + canonical_to_kind: dict[str, str] + canonical_to_description: dict[str, str | None] = {} def _collect_render_pairs( *, - memories: List[Memory], - visible_models: List[SlayerModel], - datasources: List[str], - datasource_descriptions: Optional[Dict[str, Optional[str]]] = None, -) -> List[Tuple[str, str, str, Optional[str]]]: + memories: list[Memory], + visible_models: list[SlayerModel], + datasources: list[str], + datasource_descriptions: dict[str, str | None] | None = None, +) -> list[tuple[str, str, str, str | None]]: """Return ``[(canonical_id, kind, rendered_text, description), ...]`` for every doc that goes into the index. Routes through the unified dispatch helpers in ``slayer.search.render`` (DEV-1513). Hidden @@ -146,8 +145,8 @@ def _collect_render_pairs( ``canonical_id → description`` map symmetrical to the existing canonical-text map. """ - out: List[Tuple[str, str, str, Optional[str]]] = [] - models_by_ds: Dict[str, List[SlayerModel]] = {} + out: list[tuple[str, str, str, str | None]] = [] + models_by_ds: dict[str, list[SlayerModel]] = {} for m in visible_models: models_by_ds.setdefault(m.data_source, []).append(m) descriptions = datasource_descriptions or {} @@ -175,10 +174,10 @@ def _collect_render_pairs( def build_in_memory_corpus( *, - memories: List[Memory], - models: List[SlayerModel], - datasources: List[str], - datasource_descriptions: Optional[Dict[str, Optional[str]]] = None, + memories: list[Memory], + models: list[SlayerModel], + datasources: list[str], + datasource_descriptions: dict[str, str | None] | None = None, ) -> Corpus: """Build the index AND the parallel canonical lookup maps in one walk. @@ -210,9 +209,9 @@ def build_in_memory_corpus( datasources=datasources, datasource_descriptions=datasource_descriptions, ) - canonical_to_text: Dict[str, str] = {} - canonical_to_kind: Dict[str, str] = {} - canonical_to_description: Dict[str, Optional[str]] = {} + canonical_to_text: dict[str, str] = {} + canonical_to_kind: dict[str, str] = {} + canonical_to_description: dict[str, str | None] = {} for canonical, kind, text, description in pairs: # Memory docs use ``id="memory:"`` and ``canonical=""`` # to match the DEV-1375 tantivy schema; entity docs use the same @@ -251,8 +250,8 @@ def _apply_kind_filter( *, query: "tantivy.Query", schema: "tantivy.Schema", - kind_filter: Optional[str], - exclude_kind: Optional[str], + kind_filter: str | None, + exclude_kind: str | None, ) -> "tantivy.Query": """Wrap ``query`` in a boolean query that ``Must`` includes (or ``MustNot`` excludes) docs whose ``kind`` field exactly equals the @@ -277,10 +276,10 @@ def search_index( index: tantivy.Index, question: str, limit: int = 20, - fields: Optional[List[str]] = None, - kind_filter: Optional[str] = None, - exclude_kind: Optional[str] = None, -) -> List[IndexHit]: + fields: list[str] | None = None, + kind_filter: str | None = None, + exclude_kind: str | None = None, +) -> list[IndexHit]: """Run a tantivy query against ``index``. Args: @@ -323,12 +322,12 @@ def search_index( ) searcher = index.searcher() raw_hits = searcher.search(query, limit).hits - out: List[IndexHit] = [] + out: list[IndexHit] = [] for score, address in raw_hits: doc = searcher.doc(address) kind = str(doc.get_first("kind")) canonical = str(doc.get_first("canonical")) - memory_id: Optional[str] = None + memory_id: str | None = None if kind == "memory": memory_id = canonical or None out.append(IndexHit( diff --git a/slayer/search/render.py b/slayer/search/render.py index 434d729b..252086bb 100644 --- a/slayer/search/render.py +++ b/slayer/search/render.py @@ -27,7 +27,6 @@ from __future__ import annotations import json -from typing import List, Optional from pydantic import BaseModel @@ -57,10 +56,10 @@ class RenderedEntity(BaseModel): canonical_id: str kind: str text: str - description: Optional[str] = None + description: str | None = None -def _named_children_csv(items: List[tuple[str, str]]) -> str: +def _named_children_csv(items: list[tuple[str, str]]) -> str: """Render ``[("a", "column"), ("b", "column")]`` as ``"a (column), b (column)"``.""" return ", ".join(f"{name} ({kind})" for name, kind in items) @@ -73,8 +72,8 @@ def _named_children_csv(items: List[tuple[str, str]]) -> str: def render_datasource_text( *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, + models: list[SlayerModel], + description: str | None = None, ) -> str: """Datasource doc: name + own description (when set) + named-child mentions for each model. @@ -87,7 +86,7 @@ def render_datasource_text( ``DatasourceConfig.description`` field, in parity with the other entity render helpers. """ - lines: List[str] = [f"Datasource: {name}"] + lines: list[str] = [f"Datasource: {name}"] if description: lines.append(f"Description: {description}") visible = [m for m in models if not m.hidden] @@ -103,8 +102,8 @@ def render_datasource_text( def render_datasource_pair( *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, + models: list[SlayerModel], + description: str | None = None, ) -> RenderedEntity: """Unified dispatch (DEV-1513) for the datasource doc. Used by both the tantivy corpus builder and the embedding refresh path so the @@ -133,7 +132,7 @@ def render_datasource_pair( def render_model_text(*, model: SlayerModel) -> str: """Model doc: own metadata, non-named children in full, named children by name + kind only.""" - lines: List[str] = [ + lines: list[str] = [ f"Model: {model.data_source}.{model.name}", ] if model.description: @@ -198,7 +197,7 @@ def render_model_text(*, model: SlayerModel) -> str: def render_column_text(*, model: SlayerModel, column: Column) -> str: """Column doc: parent qualifier + per-field metadata + cached sample.""" - lines: List[str] = [ + lines: list[str] = [ f"Column: {model.data_source}.{model.name}.{column.name}", f"Type: {column.type}", ] @@ -257,7 +256,7 @@ def render_column_text(*, model: SlayerModel, column: Column) -> str: def render_measure_text(*, model: SlayerModel, measure: ModelMeasure) -> str: name = measure.name or "" - lines: List[str] = [ + lines: list[str] = [ f"Measure: {model.data_source}.{model.name}.{name}", f"Formula: {measure.formula}", ] @@ -274,7 +273,7 @@ def render_measure_text(*, model: SlayerModel, measure: ModelMeasure) -> str: def render_aggregation_text(*, model: SlayerModel, aggregation: Aggregation) -> str: - lines: List[str] = [ + lines: list[str] = [ f"Aggregation: {model.data_source}.{model.name}.{aggregation.name}", ] if aggregation.formula: @@ -303,7 +302,7 @@ def render_memory_text(*, memory: Memory) -> str: ``Memory.description``. Without this, installs without the optional embedding extra would lose recall for the new field. """ - lines: List[str] = [memory.learning] + lines: list[str] = [memory.learning] if memory.description: lines.append(memory.description) if memory.entities: @@ -319,7 +318,7 @@ def compact_description_from_learning(learning: str) -> str: No special-case for ``description:`` keyword lines (the user explicitly rejected that during spec review). """ - para: List[str] = [] + para: list[str] = [] started = False for line in learning.splitlines(): if line.strip(): @@ -352,7 +351,7 @@ def render_memory_text_for_embedding(*, memory: Memory) -> str: # --------------------------------------------------------------------------- -def collect_model_entity_pairs(*, model: SlayerModel) -> List[RenderedEntity]: +def collect_model_entity_pairs(*, model: SlayerModel) -> list[RenderedEntity]: """Walk a model's subtree (model + visible columns + named measures + custom aggregations) into the unified ``RenderedEntity`` shape. @@ -372,7 +371,7 @@ def collect_model_entity_pairs(*, model: SlayerModel) -> List[RenderedEntity]: if model.hidden: return [] qualifier = f"{model.data_source}.{model.name}" - out: List[RenderedEntity] = [RenderedEntity( + out: list[RenderedEntity] = [RenderedEntity( canonical_id=qualifier, kind="model", text=render_model_text(model=model), diff --git a/slayer/search/retriever.py b/slayer/search/retriever.py index baab5f54..b0799853 100644 --- a/slayer/search/retriever.py +++ b/slayer/search/retriever.py @@ -26,7 +26,6 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Dict, List, Optional from pydantic import BaseModel, Field @@ -51,10 +50,10 @@ class RetrievalResult(BaseModel): query embedding failed"). """ - memory_ranking: List[str] = Field(default_factory=list) - text_by_id: Dict[str, str] = Field(default_factory=dict) - entity_ranking: List[str] = Field(default_factory=list) - warnings: List[str] = Field(default_factory=list) + memory_ranking: list[str] = Field(default_factory=list) + text_by_id: dict[str, str] = Field(default_factory=dict) + entity_ranking: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) class Retriever(ABC): @@ -68,12 +67,12 @@ class Retriever(ABC): async def retrieve( self, *, - query_entities: List[str], - question: Optional[str], - all_memories: List[Memory], + query_entities: list[str], + question: str | None, + all_memories: list[Memory], valid_canonicals: set, - corpus: Optional[Corpus], - datasource: Optional[str], + corpus: Corpus | None, + datasource: str | None, ) -> RetrievalResult: """Return the channel's memory + entity rankings (plus any per-call warnings). Empty rankings indicate "channel inactive @@ -84,19 +83,19 @@ async def retrieve( # Create / refresh hooks — default no-op # ------------------------------------------------------------------ - async def upsert_memory(self, memory: Memory) -> List[str]: # NOSONAR(S7503) — async signature required by Retriever ABC; subclasses override with truly-async hooks + async def upsert_memory(self, memory: Memory) -> list[str]: # NOSONAR(S7503) — async signature required by Retriever ABC; subclasses override with truly-async hooks return [] - async def refresh_model_subtree(self, model: SlayerModel) -> List[str]: # NOSONAR(S7503) — async signature required by Retriever ABC; subclasses override with truly-async hooks + async def refresh_model_subtree(self, model: SlayerModel) -> list[str]: # NOSONAR(S7503) — async signature required by Retriever ABC; subclasses override with truly-async hooks return [] async def refresh_datasource( # NOSONAR(S7503) — async signature required by Retriever ABC; subclasses override with truly-async hooks self, *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, - ) -> List[str]: + models: list[SlayerModel], + description: str | None = None, + ) -> list[str]: return [] # ------------------------------------------------------------------ diff --git a/slayer/search/retrievers/bm25.py b/slayer/search/retrievers/bm25.py index 773c790f..597df6db 100644 --- a/slayer/search/retrievers/bm25.py +++ b/slayer/search/retrievers/bm25.py @@ -13,7 +13,6 @@ from __future__ import annotations -from typing import List, Optional from slayer.memories.models import MEMORY_CANONICAL_PREFIX as _MEMORY_PREFIX from slayer.memories.models import Memory @@ -23,12 +22,12 @@ def _filter_memories_entities( - memories: List[Memory], *, valid_canonicals: set, -) -> List[Memory]: + memories: list[Memory], *, valid_canonicals: set, +) -> list[Memory]: """Return shallow copies of ``memories`` whose ``entities`` lists are filtered down to ``valid_canonicals`` only (DEV-1428). Used to feed BM25 a stale-free corpus without writing back to storage.""" - out: List[Memory] = [] + out: list[Memory] = [] for m in memories: live = [e for e in m.entities if e in valid_canonicals] if live == m.entities: @@ -38,11 +37,11 @@ def _filter_memories_entities( return out -def _augment_with_self_refs(memories: List[Memory]) -> List[Memory]: +def _augment_with_self_refs(memories: list[Memory]) -> list[Memory]: """DEV-1513: augment each memory's ``entities`` with ``memory:`` so a user-supplied ``memory:`` ref surfaces the named memory at the top of the BM25 ranking. Idempotent.""" - out: List[Memory] = [] + out: list[Memory] = [] for m in memories: self_ref = f"{_MEMORY_PREFIX}{m.id}" if self_ref in m.entities: @@ -63,12 +62,12 @@ class BM25Retriever(Retriever): async def retrieve( self, *, - query_entities: List[str], - question: Optional[str], - all_memories: List[Memory], + query_entities: list[str], + question: str | None, + all_memories: list[Memory], valid_canonicals: set, - corpus: Optional[Corpus], - datasource: Optional[str], + corpus: Corpus | None, + datasource: str | None, ) -> RetrievalResult: if not query_entities: return RetrievalResult() diff --git a/slayer/search/retrievers/embeddings.py b/slayer/search/retrievers/embeddings.py index d429b481..44000031 100644 --- a/slayer/search/retrievers/embeddings.py +++ b/slayer/search/retrievers/embeddings.py @@ -11,7 +11,6 @@ import hashlib import logging -from typing import List, Optional, Set, Tuple from slayer.core.models import SlayerModel from slayer.embeddings import client as embedding_client @@ -46,7 +45,7 @@ def _memory_canonical_id(memory_id: str) -> str: return f"{_MEMORY_PREFIX}{memory_id}" -def _memory_id_from_canonical(canonical_id: str) -> Optional[str]: +def _memory_id_from_canonical(canonical_id: str) -> str | None: """Parse a memory row's canonical id back into the str memory id. Returns ``None`` when the input is not exactly of the shape @@ -65,11 +64,11 @@ def _memory_id_from_canonical(canonical_id: str) -> Optional[str]: def _filter_embedding_corpus_by_datasource( - rows: List[Embedding], + rows: list[Embedding], *, datasource: str, - eligible_memory_canonicals: Set[str], -) -> List[Embedding]: + eligible_memory_canonicals: set[str], +) -> list[Embedding]: """DEV-1409: narrow the embedding corpus to rows that survive a datasource filter. Memory rows must appear in ``eligible_memory_canonicals`` (already datasource-filtered @@ -90,7 +89,7 @@ def _filter_embedding_corpus_by_datasource( def _rank_embedding_kind( *, rows, normalised_query, np, normalise_matrix, top_k_cosine, -) -> List[str]: +) -> list[str]: """Rank one kind of embedding rows by cosine similarity to the pre-normalised query vector. Returns the rows' ``canonical_id`` strings in descending similarity order.""" @@ -153,7 +152,7 @@ def __init__( self, *, storage: StorageBackend, - model_name: Optional[str] = None, + model_name: str | None = None, ) -> None: self._storage = storage self._model_name = model_name or current_model() @@ -166,13 +165,13 @@ def model_name(self) -> str: # Read side # ------------------------------------------------------------------ - async def fetch_corpus(self) -> List[Embedding]: + async def fetch_corpus(self) -> list[Embedding]: """Return every embedding row under the active model name.""" return await self._storage.list_embeddings( embedding_model_name=self._model_name, ) - async def embed_question(self, question: str) -> Optional[List[float]]: + async def embed_question(self, question: str) -> list[float] | None: """Embed a search query string. ``None`` when the channel is unavailable or the call fails.""" return await embedding_client.embed_query( @@ -182,12 +181,12 @@ async def embed_question(self, question: str) -> Optional[List[float]]: async def retrieve( self, *, - query_entities: List[str], - question: Optional[str], - all_memories: List[Memory], + query_entities: list[str], + question: str | None, + all_memories: list[Memory], valid_canonicals: set, - corpus: Optional[Corpus], - datasource: Optional[str], + corpus: Corpus | None, + datasource: str | None, ) -> RetrievalResult: """Run cosine over the embedding corpus, returning BOTH memory and entity rankings from a single ``fetch_corpus`` + @@ -271,7 +270,7 @@ async def retrieve( normalise_matrix=normalise_matrix, top_k_cosine=top_k_cosine, ) - memory_ranking: List[str] = [] + memory_ranking: list[str] = [] for canonical in ranked_memory_canonicals: memory_id = _memory_id_from_canonical(canonical) if memory_id is not None: @@ -292,7 +291,7 @@ async def retrieve( # Write side — create / refresh hooks # ------------------------------------------------------------------ - async def upsert_memory(self, memory: Memory) -> List[str]: + async def upsert_memory(self, memory: Memory) -> list[str]: """Refresh the embedding for a single memory. Returns warning strings (empty on success or hash-skip). Stays silent on the write path when the channel is unavailable — that's "feature @@ -312,9 +311,9 @@ async def refresh_datasource( self, *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, - ) -> List[str]: + models: list[SlayerModel], + description: str | None = None, + ) -> list[str]: """Refresh the embedding for one datasource doc. Routes through the unified :func:`render_datasource_pair` so the @@ -337,7 +336,7 @@ async def refresh_datasource( ) return await self._apply_pending([pending]) - async def refresh_model_subtree(self, model: SlayerModel) -> List[str]: + async def refresh_model_subtree(self, model: SlayerModel) -> list[str]: """Refresh the model doc + every visible column + named measures + custom aggregations in a single batch call. @@ -348,7 +347,7 @@ async def refresh_model_subtree(self, model: SlayerModel) -> List[str]: """ if not embedding_client.is_available(): return [] - pending: List[_PendingRefresh] = [ + pending: list[_PendingRefresh] = [ _PendingRefresh( canonical_id=re.canonical_id, entity_kind=re.kind, # type: ignore[arg-type] @@ -363,8 +362,8 @@ async def refresh_model_subtree(self, model: SlayerModel) -> List[str]: # ------------------------------------------------------------------ async def _apply_pending( - self, pending: List[_PendingRefresh], - ) -> List[str]: + self, pending: list[_PendingRefresh], + ) -> list[str]: """Hash-skip, batch-embed, and persist. Returns warning strings. DEV-1405: two batched storage round-trips per call — one @@ -378,8 +377,8 @@ async def _apply_pending( return [] texts = [p.text for p in stale] vectors = await embed_batch(texts, model=self._model_name) - warnings: List[str] = [] - rows: List[Embedding] = [] + warnings: list[str] = [] + rows: list[Embedding] = [] for p, vec in zip(stale, vectors): if vec is None: warnings.append( @@ -411,8 +410,8 @@ async def _apply_pending( return warnings async def _filter_stale( - self, pending: List[_PendingRefresh], - ) -> Tuple[List[_PendingRefresh], int]: + self, pending: list[_PendingRefresh], + ) -> tuple[list[_PendingRefresh], int]: """Drop pending entries whose stored content_hash already matches. Returns ``(stale_entries, fresh_skipped_count)``. DEV-1405: one batched ``get_embeddings_for_canonical_ids`` call @@ -421,7 +420,7 @@ async def _filter_stale( canonical_ids=[p.canonical_id for p in pending], embedding_model_name=self._model_name, ) - stale: List[_PendingRefresh] = [] + stale: list[_PendingRefresh] = [] fresh = 0 for p in pending: match = existing.get(p.canonical_id) diff --git a/slayer/search/retrievers/tantivy.py b/slayer/search/retrievers/tantivy.py index 0fd30b68..f4efed4f 100644 --- a/slayer/search/retrievers/tantivy.py +++ b/slayer/search/retrievers/tantivy.py @@ -14,14 +14,13 @@ from __future__ import annotations -from typing import Dict, List, Optional, Tuple from slayer.memories.models import Memory from slayer.search.index import Corpus, IndexHit, search_index from slayer.search.retriever import RetrievalResult, Retriever -def _count_corpus_kinds(corpus: Corpus) -> Tuple[int, int]: +def _count_corpus_kinds(corpus: Corpus) -> tuple[int, int]: """Return ``(memory_count, entity_count)`` for a built corpus. Used to pass ``limit = full per-kind corpus size`` to each @@ -47,18 +46,18 @@ class TantivyRetriever(Retriever): async def retrieve( self, *, - query_entities: List[str], - question: Optional[str], - all_memories: List[Memory], + query_entities: list[str], + question: str | None, + all_memories: list[Memory], valid_canonicals: set, - corpus: Optional[Corpus], - datasource: Optional[str], + corpus: Corpus | None, + datasource: str | None, ) -> RetrievalResult: if corpus is None or not question or not question.strip(): return RetrievalResult() memory_count, entity_count = _count_corpus_kinds(corpus) - memory_hits: List[IndexHit] = ( + memory_hits: list[IndexHit] = ( search_index( index=corpus.index, question=question, @@ -68,7 +67,7 @@ async def retrieve( if memory_count > 0 else [] ) - entity_hits: List[IndexHit] = ( + entity_hits: list[IndexHit] = ( search_index( index=corpus.index, question=question, @@ -79,8 +78,8 @@ async def retrieve( else [] ) - memory_ranking: List[str] = [] - text_by_id: Dict[str, str] = {} + memory_ranking: list[str] = [] + text_by_id: dict[str, str] = {} for hit in memory_hits: if hit.memory_id is None: continue diff --git a/slayer/search/rrf.py b/slayer/search/rrf.py index f4dbe666..b243396a 100644 --- a/slayer/search/rrf.py +++ b/slayer/search/rrf.py @@ -13,14 +13,15 @@ from __future__ import annotations -from typing import Hashable, List, TypeVar +from typing import TypeVar +from collections.abc import Hashable K = TypeVar("K", bound=Hashable) def rrf_fuse( *, - rankings: List[List[K]], + rankings: list[list[K]], k: int = 60, ) -> dict[K, float]: """Fuse multiple ranked lists into one score map. diff --git a/slayer/search/service.py b/slayer/search/service.py index 94477eb4..ed22da0e 100644 --- a/slayer/search/service.py +++ b/slayer/search/service.py @@ -68,7 +68,7 @@ import asyncio import logging -from typing import Any, Dict, FrozenSet, List, Optional, Set, Tuple, Union +from typing import Any from pydantic import BaseModel, Field @@ -140,9 +140,9 @@ class SearchHit(BaseModel): id: str score: float text: str - description: Optional[str] = None - matched_entities: List[str] = Field(default_factory=list) - query: Optional[SlayerQuery] = None + description: str | None = None + matched_entities: list[str] = Field(default_factory=list) + query: SlayerQuery | None = None # --------------------------------------------------------------------------- @@ -158,7 +158,7 @@ class LookupFound(BaseModel): kind: str text: str - description: Optional[str] = None + description: str | None = None class LookupHidden(BaseModel): @@ -177,7 +177,7 @@ class LookupMissing(BaseModel): pass -LookupResult = Union[LookupFound, LookupHidden, LookupMissing] +LookupResult = LookupFound | LookupHidden | LookupMissing class SearchResponse(BaseModel): @@ -185,9 +185,9 @@ class SearchResponse(BaseModel): list ranked by RRF score; consumers partition by ``kind`` (or by ``query is None`` for the memory subset) at the call site.""" - results: List[SearchHit] = Field(default_factory=list) - resolved_input_entities: List[str] = Field(default_factory=list) - warnings: List[str] = Field(default_factory=list) + results: list[SearchHit] = Field(default_factory=list) + resolved_input_entities: list[str] = Field(default_factory=list) + warnings: list[str] = Field(default_factory=list) # --------------------------------------------------------------------------- @@ -195,7 +195,7 @@ class SearchResponse(BaseModel): # --------------------------------------------------------------------------- -def _coerce_query(query: Union[SlayerQuery, dict]) -> SlayerQuery: +def _coerce_query(query: SlayerQuery | dict) -> SlayerQuery: if isinstance(query, SlayerQuery): return query if isinstance(query, dict): @@ -205,9 +205,9 @@ def _coerce_query(query: Union[SlayerQuery, dict]) -> SlayerQuery: ) -def _dedup(items: List[str]) -> List[str]: - seen: Set[str] = set() - out: List[str] = [] +def _dedup(items: list[str]) -> list[str]: + seen: set[str] = set() + out: list[str] = [] for x in items: if x not in seen: seen.add(x) @@ -216,8 +216,8 @@ def _dedup(items: List[str]) -> List[str]: def _filter_memories_by_datasource( - *, memories: List[Memory], datasource: Optional[str], -) -> List[Memory]: + *, memories: list[Memory], datasource: str | None, +) -> list[Memory]: """DEV-1409: keep memories with at least one entity rooted at ``datasource``. ``datasource=None`` is a no-op identity filter so callers can call this unconditionally.""" @@ -232,7 +232,7 @@ def _filter_memories_by_datasource( ] -def _collect_memory_canonicals(memories: List[Memory]) -> set: +def _collect_memory_canonicals(memories: list[Memory]) -> set: return {f"{_MEMORY_PREFIX}{m.id}" for m in memories} @@ -260,9 +260,9 @@ def _build_memory_hit( mem: Memory, memory_id: str, score: float, - text_by_id: Dict[str, str], - canonical_input_entities: List[str], - valid_canonicals: Optional[set] = None, + text_by_id: dict[str, str], + canonical_input_entities: list[str], + valid_canonicals: set | None = None, compact: bool = True, ) -> SearchHit: """Build a SearchHit for a memory (DEV-1532 unified shape). @@ -318,9 +318,9 @@ def _build_memory_hit( def _resolve_entity_hit_kind_text( *, canonical: str, - corpus: Optional[Corpus], - named_kind_text: Optional[Dict[str, Tuple[str, str, Optional[str]]]], -) -> Optional[Tuple[str, str, Optional[str]]]: + corpus: Corpus | None, + named_kind_text: dict[str, tuple[str, str, str | None]] | None, +) -> tuple[str, str, str | None] | None: """DEV-1513 / DEV-1549: resolve one canonical's ``(kind, text, description)`` triple for an entity hit. Prefers the corpus (channels 2/3 already built it); falls back to the channel-1 @@ -345,15 +345,15 @@ def _build_hit_from_fused_key( key: str, score: float, memory_by_id: dict, - text_by_id: Dict[str, str], - canonical_input_entities: List[str], - corpus: Optional[Corpus], - named_kind_text: Optional[Dict[str, Tuple[str, str, Optional[str]]]], - valid_canonicals: Optional[set], - candidate_ids: Optional[FrozenSet[str]], - kind_filter: Optional[Set[str]], + text_by_id: dict[str, str], + canonical_input_entities: list[str], + corpus: Corpus | None, + named_kind_text: dict[str, tuple[str, str, str | None]] | None, + valid_canonicals: set | None, + candidate_ids: frozenset[str] | None, + kind_filter: set[str] | None, compact: bool = True, -) -> Optional[SearchHit]: +) -> SearchHit | None: """Build one SearchHit from a fused (key, score) pair, or return None to skip. Applies the DEV-1464 cypher_filter (candidate_ids allowlist for the full graph path; kind_filter for the naive fallback) BEFORE @@ -405,19 +405,19 @@ def _build_hit_from_fused_key( def _fuse_all_hits( *, - memory_rankings: List[List[str]], - entity_rankings: List[List[str]], + memory_rankings: list[list[str]], + entity_rankings: list[list[str]], memory_by_id: dict, - text_by_id: Dict[str, str], - canonical_input_entities: List[str], - corpus: Optional[Corpus], - named_kind_text: Optional[Dict[str, Tuple[str, str, Optional[str]]]], + text_by_id: dict[str, str], + canonical_input_entities: list[str], + corpus: Corpus | None, + named_kind_text: dict[str, tuple[str, str, str | None]] | None, max_results: int, - valid_canonicals: Optional[set] = None, - candidate_ids: Optional[FrozenSet[str]] = None, - kind_filter: Optional[Set[str]] = None, + valid_canonicals: set | None = None, + candidate_ids: frozenset[str] | None = None, + kind_filter: set[str] | None = None, compact: bool = True, -) -> List[SearchHit]: +) -> list[SearchHit]: """RRF-fuse memory and entity rankings into a single flat list (DEV-1532). Memory IDs are prefixed with the canonical memory prefix so the unified pool contains no key collisions. @@ -437,7 +437,7 @@ def _fuse_all_hits( fused = rrf_fuse(rankings=non_empty, k=_RRF_K) if non_empty else {} fused_sorted = sorted(fused.items(), key=lambda kv: kv[1], reverse=True) - results: List[SearchHit] = [] + results: list[SearchHit] = [] for key, score in fused_sorted: hit = _build_hit_from_fused_key( key=key, @@ -460,11 +460,11 @@ def _fuse_all_hits( def _merge_text_by_id_in_declaration_order( - results: List[RetrievalResult], -) -> Dict[str, str]: + results: list[RetrievalResult], +) -> dict[str, str]: """Merge ``text_by_id`` across retriever results. First-non-empty in retriever declaration order wins per memory id.""" - merged: Dict[str, str] = {} + merged: dict[str, str] = {} for result in results: for mem_id, text in result.text_by_id.items(): if mem_id not in merged and text: @@ -479,10 +479,10 @@ def _merge_text_by_id_in_declaration_order( def _memory_id_off_datasource_warnings( *, - canonical_input_entities: List[str], - live_memory_ids: Set[str], - datasource: Optional[str], -) -> List[str]: + canonical_input_entities: list[str], + live_memory_ids: set[str], + datasource: str | None, +) -> list[str]: """DEV-1513: emit one warning per user-supplied ``memory:`` ref whose memory was dropped by the datasource pre-filter (the memory has no entities rooted at ``datasource``). Mirrors the entity-side @@ -491,7 +491,7 @@ def _memory_id_off_datasource_warnings( No-op when ``datasource`` is None (nothing was filtered out).""" if datasource is None: return [] - out: List[str] = [] + out: list[str] = [] for canonical in canonical_input_entities: if not canonical.startswith(_MEMORY_PREFIX): continue @@ -506,9 +506,9 @@ def _memory_id_off_datasource_warnings( def _memory_id_cypher_filter_warnings( *, - canonical_input_entities: List[str], - candidate_ids: FrozenSet[str], -) -> List[str]: + canonical_input_entities: list[str], + candidate_ids: frozenset[str], +) -> list[str]: """DEV-1464: emit one warning per user-supplied ``memory:`` ref that was excluded by the cypher_filter allowlist (the graph query did not return that memory's canonical id).""" @@ -529,7 +529,7 @@ async def _lookup_bare_datasource_canonical( if ds not in known: return LookupMissing() identities = await storage._list_all_model_identities() - models: List[SlayerModel] = [] + models: list[SlayerModel] = [] for ident_ds, name in identities: if ident_ds != ds: continue @@ -551,7 +551,7 @@ async def _lookup_model_or_leaf_canonical( canonical: str, ds: str, model_name: str, - leaf: Optional[str], + leaf: str | None, storage: StorageBackend, ) -> LookupResult: """DEV-1513: ``.`` and ``..`` branches of @@ -579,7 +579,7 @@ async def _lookup_named_entity( *, canonical: str, storage: StorageBackend, - corpus: Optional[Corpus], + corpus: Corpus | None, ) -> LookupResult: """Resolve a canonical id to its ``(kind, text, description)`` triple for channel-1 named-entity surfacing (DEV-1513 / DEV-1549).""" @@ -611,8 +611,8 @@ async def _lookup_named_entity( def _group_column_hits( - results: List[SearchHit], -) -> Dict[Tuple[str, str], List[Tuple[int, SearchHit, str]]]: + results: list[SearchHit], +) -> dict[tuple[str, str], list[tuple[int, SearchHit, str]]]: """DEV-1516 helper: split a fused result list into per-model buckets for the search-side sample-refresh hook. @@ -623,7 +623,7 @@ def _group_column_hits( models. Each member tuple is ``(original_hit_index, hit, column_name)`` — the index is preserved so caller can splice refreshed text back into the original list in place.""" - groups: Dict[Tuple[str, str], List[Tuple[int, SearchHit, str]]] = {} + groups: dict[tuple[str, str], list[tuple[int, SearchHit, str]]] = {} for idx, hit in enumerate(results): if hit.kind != "column": continue @@ -649,8 +649,8 @@ def __init__( self, *, storage: StorageBackend, - engine: Optional[SlayerQueryEngine] = None, - retrievers: Optional[List[Retriever]] = None, + engine: SlayerQueryEngine | None = None, + retrievers: list[Retriever] | None = None, ) -> None: """DEV-1516: ``engine`` is optional so storage-only test contexts keep working unchanged. When supplied, the post-fusion column-hit @@ -659,13 +659,13 @@ def __init__( Without an engine the hook is a silent no-op.""" self._storage = storage self._engine = engine - self._retrievers: List[Retriever] = ( + self._retrievers: list[Retriever] = ( list(retrievers) if retrievers is not None else self._default_retrievers(storage) ) @staticmethod - def _default_retrievers(storage: StorageBackend) -> List[Retriever]: + def _default_retrievers(storage: StorageBackend) -> list[Retriever]: return [ BM25Retriever(), TantivyRetriever(), @@ -673,15 +673,15 @@ def _default_retrievers(storage: StorageBackend) -> List[Retriever]: ] @property - def retrievers(self) -> List[Retriever]: + def retrievers(self) -> list[Retriever]: return self._retrievers async def _refresh_stale_column_hits( self, *, - results: List[SearchHit], + results: list[SearchHit], compact: bool = True, - ) -> List[SearchHit]: + ) -> list[SearchHit]: """DEV-1516 post-fusion column-hit refresh. Groups column hits by ``(data_source, model_name)`` and dispatches @@ -700,7 +700,7 @@ async def _refresh_stale_column_hits( groups = _group_column_hits(results) if not groups: return results - refreshed_by_idx: Dict[int, SearchHit] = {} + refreshed_by_idx: dict[int, SearchHit] = {} await asyncio.gather(*[ self._refresh_group_worker( ds_name=ds, model_name=model_name, @@ -720,8 +720,8 @@ async def _refresh_group_worker( *, ds_name: str, model_name: str, - members: List[Tuple[int, SearchHit, str]], - refreshed_by_idx: Dict[int, SearchHit], + members: list[tuple[int, SearchHit, str]], + refreshed_by_idx: dict[int, SearchHit], compact: bool = True, ) -> None: """Refresh every column hit on one ``(data_source, model_name)`` @@ -758,7 +758,7 @@ async def _refresh_group_worker( # Helper returned the input — cache hit, ineligible, or # any failure. Leave the hit text as-is. continue - update: Dict[str, Any] = { + update: dict[str, Any] = { "description": refreshed_col.description, } if not compact: @@ -774,11 +774,11 @@ async def _refresh_group_worker( async def search( # NOSONAR(S3776) — single orchestrator entry point; stages are linear and named self, *, - entities: Optional[List[str]] = None, - query: Optional[Union[SlayerQuery, dict]] = None, - question: Optional[str] = None, - datasource: Optional[str] = None, - cypher_filter: Optional[str] = None, + entities: list[str] | None = None, + query: SlayerQuery | dict | None = None, + question: str | None = None, + datasource: str | None = None, + cypher_filter: str | None = None, max_results: int = 10, compact: bool = True, ) -> SearchResponse: @@ -831,7 +831,7 @@ async def search( # NOSONAR(S3776) — single orchestrator entry point; stages # Datasource filter runs first so the off-datasource warning # reflects "memory dropped because of datasource", not # "memory dropped because of cypher_filter". - datasource_filtered_memories: List[Memory] = ( + datasource_filtered_memories: list[Memory] = ( _filter_memories_by_datasource( memories=await self._storage.list_memories(entities=None), datasource=datasource, @@ -854,7 +854,7 @@ async def search( # NOSONAR(S3776) — single orchestrator entry point; stages # retrieval path — BM25 / tantivy / embeddings rank only the # surviving memories. if candidate_ids is not None: - all_memories: List[Memory] = [ + all_memories: list[Memory] = [ m for m in datasource_filtered_memories if f"{_MEMORY_PREFIX}{m.id}" in candidate_ids ] @@ -865,7 +865,7 @@ async def search( # NOSONAR(S3776) — single orchestrator entry point; stages all_memories=all_memories, datasource=datasource, ) - corpus: Optional[Corpus] = None + corpus: Corpus | None = None if question_active: all_models, datasources, datasource_descriptions = ( await self._collect_index_corpus(datasource=datasource) @@ -917,7 +917,7 @@ async def search( # NOSONAR(S3776) — single orchestrator entry point; stages ), return_exceptions=True, ) - results: List[RetrievalResult] = [] + results: list[RetrievalResult] = [] for r, raw in zip(self._retrievers, raw_results): if isinstance(raw, BaseException): warnings.append( @@ -994,13 +994,13 @@ async def search( # NOSONAR(S3776) — single orchestrator entry point; stages async def _apply_cypher_filter( self, *, - cypher_filter: Optional[str], - canonical_input_entities: List[str], - warnings: List[str], - ) -> Tuple[ - Optional[FrozenSet[str]], - Optional[Set[str]], - Optional[SearchResponse], + cypher_filter: str | None, + canonical_input_entities: list[str], + warnings: list[str], + ) -> tuple[ + frozenset[str] | None, + set[str] | None, + SearchResponse | None, ]: """DEV-1464: resolve the optional ``cypher_filter`` into ``(candidate_ids, kind_filter, early)``. @@ -1038,11 +1038,11 @@ async def _apply_cypher_filter( async def _build_channel_1_entity_ranking( self, *, - canonical_input_entities: List[str], - datasource: Optional[str], - corpus: Optional[Corpus], - candidate_ids: Optional[FrozenSet[str]] = None, - ) -> Tuple[List[str], Dict[str, Tuple[str, str, Optional[str]]], List[str]]: + canonical_input_entities: list[str], + datasource: str | None, + corpus: Corpus | None, + candidate_ids: frozenset[str] | None = None, + ) -> tuple[list[str], dict[str, tuple[str, str, str | None]], list[str]]: """DEV-1513: produce channel-1's contribution to the entity ranking by surfacing each user-named canonical ref as itself. @@ -1062,9 +1062,9 @@ async def _build_channel_1_entity_ranking( the allowlist are dropped (with a warning) BEFORE the rendering / lookup work so we don't waste a storage round-trip. """ - entity_ranking: List[str] = [] - named_kind_text: Dict[str, Tuple[str, str, Optional[str]]] = {} - warnings: List[str] = [] + entity_ranking: list[str] = [] + named_kind_text: dict[str, tuple[str, str, str | None]] = {} + warnings: list[str] = [] for canonical in canonical_input_entities: if canonical.startswith(_MEMORY_PREFIX): # memory: refs participate in the memory ranking only. @@ -1108,7 +1108,7 @@ async def _build_channel_1_entity_ranking( # Write side — fan-out to retrievers # ------------------------------------------------------------------ - async def upsert_memory(self, memory: Memory) -> List[str]: + async def upsert_memory(self, memory: Memory) -> list[str]: return await self._fan_out_with_isolation( hook_name="upsert_memory", invoke=lambda r: r.upsert_memory(memory), @@ -1116,7 +1116,7 @@ async def upsert_memory(self, memory: Memory) -> List[str]: async def refresh_model_subtree( self, model: SlayerModel, - ) -> List[str]: + ) -> list[str]: return await self._fan_out_with_isolation( hook_name="refresh_model_subtree", invoke=lambda r: r.refresh_model_subtree(model), @@ -1126,9 +1126,9 @@ async def refresh_datasource( self, *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, - ) -> List[str]: + models: list[SlayerModel], + description: str | None = None, + ) -> list[str]: return await self._fan_out_with_isolation( hook_name="refresh_datasource", invoke=lambda r: r.refresh_datasource( @@ -1138,12 +1138,12 @@ async def refresh_datasource( async def _fan_out_with_isolation( self, *, hook_name: str, invoke, - ) -> List[str]: + ) -> list[str]: """Call ``invoke(retriever)`` on every registered retriever in declaration order, isolating per-retriever exceptions as prefixed warnings so subsequent retrievers still run. Returns the deduped warning list.""" - warnings: List[str] = [] + warnings: list[str] = [] for r in self._retrievers: try: warnings.extend(await invoke(r)) @@ -1158,7 +1158,7 @@ async def _fan_out_with_isolation( # ------------------------------------------------------------------ async def _validate_datasource_known( - self, datasource: Optional[str], + self, datasource: str | None, ) -> None: """DEV-1409: reject typos in ``datasource`` before any corpus walk.""" @@ -1173,14 +1173,14 @@ async def _validate_datasource_known( async def _resolve_inputs( self, *, - entities: Optional[List[str]], - query: Optional[Union[SlayerQuery, dict]], - ) -> Tuple[List[str], List[str]]: + entities: list[str] | None, + query: SlayerQuery | dict | None, + ) -> tuple[list[str], list[str]]: """Walk ``entities`` + ``query`` into a deduped canonical-entity list plus a deduped warning list. DEV-1428: lenient — per-token failures become warnings.""" - canonical: List[str] = [] - warnings: List[str] = [] + canonical: list[str] = [] + warnings: list[str] = [] if entities: for raw in entities: if not isinstance(raw, str): @@ -1213,10 +1213,10 @@ async def _recency_fallback( self, *, max_results: int, - warnings: List[str], - datasource: Optional[str] = None, - candidate_ids: Optional[FrozenSet[str]] = None, - kind_filter: Optional[Set[str]] = None, + warnings: list[str], + datasource: str | None = None, + candidate_ids: frozenset[str] | None = None, + kind_filter: set[str] | None = None, compact: bool = True, ) -> SearchResponse: """Empty-input branch: return the newest memories (both @@ -1266,7 +1266,7 @@ async def _recency_fallback( valid_canonicals = await self._valid_canonical_set( all_memories=recency_memories, datasource=datasource, ) - hits: List[SearchHit] = [] + hits: list[SearchHit] = [] for m in recency_memories: if len(hits) >= max_results: break @@ -1297,8 +1297,8 @@ async def _recency_fallback( async def _valid_canonical_set( self, *, - all_memories: List[Memory], - datasource: Optional[str], + all_memories: list[Memory], + datasource: str | None, ) -> set: canonicals: set = set() canonicals.update( @@ -1313,7 +1313,7 @@ async def _valid_canonical_set( return canonicals async def _collect_datasource_canonicals( - self, *, datasource: Optional[str], + self, *, datasource: str | None, ) -> set: names = await self._storage.list_datasources() if datasource is not None: @@ -1321,7 +1321,7 @@ async def _collect_datasource_canonicals( return set(names) async def _collect_model_subtree_canonicals( - self, *, datasource: Optional[str], + self, *, datasource: str | None, ) -> set: out: set = set() identities = await self._storage._list_all_model_identities() @@ -1345,13 +1345,13 @@ async def _collect_model_subtree_canonicals( async def _stale_query_warnings( self, *, - query_bearing_hits: List[SearchHit], - memory_by_id: Dict[str, Memory], - ) -> List[str]: + query_bearing_hits: list[SearchHit], + memory_by_id: dict[str, Memory], + ) -> list[str]: """Emit one warning per surfaced query-bearing hit whose attached ``SlayerQuery`` no longer resolves (entities pointing at deleted/renamed models or columns). DEV-1428.""" - out: List[str] = [] + out: list[str] = [] for hit in query_bearing_hits: mem = memory_by_id.get(hit.id) if mem is None or mem.query is None: @@ -1370,17 +1370,17 @@ async def _stale_query_warnings( async def _stale_query_warnings_for_named_memory_refs( self, *, - canonical_input_entities: List[str], - all_memories: List[Memory], - already_warned_ids: Set[str], - ) -> List[str]: + canonical_input_entities: list[str], + all_memories: list[Memory], + already_warned_ids: set[str], + ) -> list[str]: """DEV-1513: emit the stale-query warning for any explicitly-named ``memory:`` ref pointing at a query-bearing memory with stale refs, regardless of whether the ``max_results`` cap suppressed the hit. The user explicitly named the memory; they deserve to know the attached query is broken.""" memories_by_id = {m.id: m for m in all_memories} - out: List[str] = [] + out: list[str] = [] for canonical in canonical_input_entities: if not canonical.startswith(_MEMORY_PREFIX): continue @@ -1404,15 +1404,15 @@ async def _stale_query_warnings_for_named_memory_refs( async def _collect_index_corpus( self, *, - datasource: Optional[str] = None, - ) -> Tuple[List[SlayerModel], List[str], Dict[str, Optional[str]]]: + datasource: str | None = None, + ) -> tuple[list[SlayerModel], list[str], dict[str, str | None]]: """DEV-1549: also returns ``{ds_name → description}`` so the corpus builder can populate ``canonical_to_description`` for datasource hits without re-loading the configs.""" datasources = await self._storage.list_datasources() if datasource is not None: datasources = [d for d in datasources if d == datasource] - models: List[SlayerModel] = [] + models: list[SlayerModel] = [] identities = await self._storage._list_all_model_identities() for ds, name in identities: if datasource is not None and ds != datasource: @@ -1420,7 +1420,7 @@ async def _collect_index_corpus( m = await self._storage.get_model(name, data_source=ds) if m is not None: models.append(m) - descriptions: Dict[str, Optional[str]] = {} + descriptions: dict[str, str | None] = {} for ds_name in datasources: cfg = await self._storage.get_datasource(ds_name) descriptions[ds_name] = cfg.description if cfg is not None else None diff --git a/slayer/sql/client.py b/slayer/sql/client.py index 1cd89d52..40e5ddb1 100644 --- a/slayer/sql/client.py +++ b/slayer/sql/client.py @@ -5,7 +5,8 @@ import functools import logging import time -from typing import Any, Awaitable, Callable, Dict, List, Optional +from typing import Any +from collections.abc import Awaitable, Callable import sqlalchemy as sa import sqlalchemy.engine.url @@ -41,7 +42,7 @@ # value or the path component of `sqlite:///:memory:` connection strings. _MEMORY_DB_NAME = ":memory:" -_sync_engines: Dict[str, sa.Engine] = {} +_sync_engines: dict[str, sa.Engine] = {} def _get_sync_engine(connection_string: str) -> sa.Engine: @@ -85,7 +86,7 @@ def _is_in_memory_sqlite(connection_string: str) -> bool: database = url.database if not database or database == _MEMORY_DB_NAME: return True - query: Dict[str, Any] = dict(url.query) if url.query else {} + query: dict[str, Any] = dict(url.query) if url.query else {} # SQLite honors `mode=memory` and the `file::memory:` URI form ONLY when # the connection is opened with URI handling enabled (`uri=true`). # Without `uri=true`, SQLite treats the database part as a literal @@ -132,7 +133,7 @@ def _register_udfs(dbapi_connection, _connection_record): def _resolve_sync_engine( connection_string: str, - override_engine: Optional[sa.Engine] = None, + override_engine: sa.Engine | None = None, ) -> sa.Engine: """Choose the engine for a sync DB call. @@ -159,7 +160,7 @@ def _get_async_engine(connection_string: str): return create_async_engine(connection_string, pool_pre_ping=True) -def _async_connection_string(connection_string: str, db_type: Optional[str]) -> Optional[str]: +def _async_connection_string(connection_string: str, db_type: str | None) -> str | None: """Convert a sync connection string to its async equivalent, or None if no async driver.""" async_scheme = _ASYNC_DRIVERS.get(db_type) if async_scheme is None: @@ -175,7 +176,7 @@ def _async_connection_string(connection_string: str, db_type: Optional[str]) -> # --------------------------------------------------------------------------- -def _map_type_code(type_code, db_type: Optional[str] = None) -> str: +def _map_type_code(type_code, db_type: str | None = None) -> str: """Map a DB-API type_code to a SLayer type category. Handles DuckDB (string type names), SQLite (Python types), @@ -234,7 +235,7 @@ class returns ``None`` so other dialects fall through to the # Postgres OIDs (from pg_type) -_PG_OID_MAP: Dict[int, str] = { +_PG_OID_MAP: dict[int, str] = { 16: "boolean", # bool 20: "number", # int8 (bigint) 21: "number", # int2 (smallint) @@ -256,7 +257,7 @@ class returns ``None`` so other dialects fall through to the } # MySQL field-type codes (aiomysql wire protocol) -_MYSQL_TYPE_MAP: Dict[int, str] = { +_MYSQL_TYPE_MAP: dict[int, str] = { 0: "number", # MYSQL_TYPE_DECIMAL 1: "boolean", # MYSQL_TYPE_TINY (TINYINT/BOOL) 2: "number", # MYSQL_TYPE_SHORT @@ -281,7 +282,7 @@ class returns ``None`` so other dialects fall through to the # ODBC SQL type codes (pyodbc with SQL Server / mssql+pyodbc driver). # Positive codes are the standard ODBC C-level SQL_* constants; negative codes # are SQL Server extensions (SQL_SS_*) defined in msodbcsql.h. -_ODBC_SQL_TYPE_MAP: Dict[int, str] = { +_ODBC_SQL_TYPE_MAP: dict[int, str] = { # Integer / numeric family 4: "number", # SQL_INTEGER 5: "number", # SQL_SMALLINT @@ -316,7 +317,7 @@ class returns ``None`` so other dialects fall through to the } -def _extract_types_from_cursor(result, db_type: Optional[str] = None) -> Dict[str, str]: +def _extract_types_from_cursor(result, db_type: str | None = None) -> dict[str, str]: """Extract {column_name: type_category} from a SQLAlchemy CursorResult. Uses cursor.description type_code when available (DuckDB, Postgres). @@ -377,7 +378,7 @@ async def _run_sync_in_thread(func, *args, **kwargs): return await loop.run_in_executor(executor, call) -def _build_type_probe_sql(sql: str, db_type: Optional[str]) -> str: +def _build_type_probe_sql(sql: str, db_type: str | None) -> str: """Build a row-limiting probe query appropriate for the target dialect.""" limit = 1 if db_type in _NEEDS_ROW_FOR_TYPES else 0 if db_type in _TSQL_DB_TYPES: @@ -385,7 +386,7 @@ def _build_type_probe_sql(sql: str, db_type: Optional[str]) -> str: return f"SELECT * FROM ({sql}) AS _types LIMIT {limit}" -def _apply_type_probe_timeout(conn, db_type: Optional[str], timeout_seconds: int) -> None: +def _apply_type_probe_timeout(conn, db_type: str | None, timeout_seconds: int) -> None: """DEV-1551: apply the dialect's statement-timeout SQL ahead of a type-probe execution. Snowflake's ``LIMIT 0`` still compiles and consumes warehouse compute, so an unbounded probe can stall on a @@ -403,7 +404,7 @@ def _apply_type_probe_timeout(conn, db_type: Optional[str], timeout_seconds: int conn.execute(sa.text(timeout_sql)) -async def _apply_type_probe_timeout_async(conn, db_type: Optional[str], timeout_seconds: int) -> None: +async def _apply_type_probe_timeout_async(conn, db_type: str | None, timeout_seconds: int) -> None: """Async sibling of ``_apply_type_probe_timeout``.""" if not db_type: return @@ -421,9 +422,9 @@ async def _apply_type_probe_timeout_async(conn, db_type: Optional[str], timeout_ def _get_column_types_sync( sql: str, connection_string: str, - db_type: Optional[str], - engine: Optional[sa.Engine] = None, -) -> Dict[str, str]: + db_type: str | None, + engine: sa.Engine | None = None, +) -> dict[str, str]: """Infer column types. Uses LIMIT 0 for cursor metadata, LIMIT 1 for SQLite. T-SQL uses SELECT TOP N instead of LIMIT.""" engine = _resolve_sync_engine(connection_string, override_engine=engine) @@ -437,8 +438,8 @@ def _get_column_types_sync( async def _get_column_types_async( sql: str, engine, - db_type: Optional[str], -) -> Dict[str, str]: + db_type: str | None, +) -> dict[str, str]: """Async version of column type inference. Uses LIMIT 0; LIMIT 1 for SQLite. T-SQL uses SELECT TOP N instead of LIMIT.""" limit_sql = _build_type_probe_sql(sql, db_type) @@ -462,7 +463,7 @@ class SlayerSQLClient: def __init__(self, datasource: DatasourceConfig): self.datasource = datasource self._async_engine = None - self._sync_engine: Optional[sa.Engine] = None + self._sync_engine: sa.Engine | None = None async def aclose(self) -> None: """Dispose the cached async engine inside the current event loop.""" @@ -491,7 +492,7 @@ def _get_async_engine(self): self._async_engine = _get_async_engine(async_conn_str) return self._async_engine - def _get_sync_engine_for_client(self) -> Optional[sa.Engine]: + def _get_sync_engine_for_client(self) -> sa.Engine | None: """Return a per-client sync engine. For ``sqlite:///:memory:`` (and equivalent URI-form variants) every @@ -519,7 +520,7 @@ async def execute( self, sql: str, timeout_seconds: int = 120, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Execute SQL asynchronously.""" async_engine = self._get_async_engine() db_type = self.datasource.type @@ -546,7 +547,7 @@ async def execute( engine=self._get_sync_engine_for_client(), ) - async def get_column_types(self, sql: str) -> Dict[str, str]: + async def get_column_types(self, sql: str) -> dict[str, str]: """Infer column types by executing SQL with LIMIT 0. Returns {column_name: type_category} where type_category is @@ -575,7 +576,7 @@ def execute_sync( self, sql: str, timeout_seconds: int = 120, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Execute SQL synchronously (for CLI, notebooks, tests).""" return _execute_with_retry_sync( sql=sql, @@ -637,11 +638,11 @@ def _is_transient_db_error(exc: BaseException) -> bool: async def _retry_with_backoff( *, sql: str, - do_call: Callable[[], Awaitable[List[Dict[str, Any]]]], + do_call: Callable[[], Awaitable[list[dict[str, Any]]]], max_attempts: int, initial_delay: float, max_delay: float, -) -> List[Dict[str, Any]]: +) -> list[dict[str, Any]]: """Retry an async DB call with exponential backoff on transient errors. `sql` is used only for the warning's excerpt so users can correlate @@ -671,12 +672,12 @@ async def _retry_with_backoff( async def _execute_with_retry_async( sql: str, engine, - db_type: Optional[str], + db_type: str | None, timeout_seconds: int = 120, max_attempts: int = 3, initial_delay: float = 1.0, max_delay: float = 10.0, -) -> List[Dict[str, Any]]: +) -> list[dict[str, Any]]: return await _retry_with_backoff( sql=sql, do_call=lambda: _execute_sql_async( @@ -691,9 +692,9 @@ async def _execute_with_retry_async( async def _execute_sql_async( sql: str, engine, - db_type: Optional[str], + db_type: str | None, timeout_seconds: int = 120, -) -> List[Dict[str, Any]]: +) -> list[dict[str, Any]]: async with engine.connect() as conn: timeout_ms = timeout_seconds * 1000 if db_type in ("mysql", "mariadb"): @@ -725,13 +726,13 @@ async def _execute_sql_async( async def _execute_with_retry_threaded( sql: str, connection_string: str, - db_type: Optional[str], + db_type: str | None, timeout_seconds: int = 120, max_attempts: int = 3, initial_delay: float = 1.0, max_delay: float = 10.0, - engine: Optional[sa.Engine] = None, -) -> List[Dict[str, Any]]: + engine: sa.Engine | None = None, +) -> list[dict[str, Any]]: return await _retry_with_backoff( sql=sql, do_call=lambda: _run_sync_in_thread( @@ -756,13 +757,13 @@ async def _execute_with_retry_threaded( def _execute_with_retry_sync( sql: str, connection_string: str, - db_type: Optional[str], + db_type: str | None, timeout_seconds: int = 120, max_attempts: int = 3, initial_delay: float = 1.0, max_delay: float = 10.0, - engine: Optional[sa.Engine] = None, -) -> List[Dict[str, Any]]: + engine: sa.Engine | None = None, +) -> list[dict[str, Any]]: delay = initial_delay for attempt in range(max_attempts): try: @@ -789,10 +790,10 @@ def _execute_with_retry_sync( def _execute_sql_sync( sql: str, connection_string: str, - db_type: Optional[str], + db_type: str | None, timeout_seconds: int = 120, - engine: Optional[sa.Engine] = None, -) -> List[Dict[str, Any]]: + engine: sa.Engine | None = None, +) -> list[dict[str, Any]]: engine = _resolve_sync_engine(connection_string, override_engine=engine) with engine.connect() as conn: timeout_ms = timeout_seconds * 1000 diff --git a/slayer/sql/dialects/__init__.py b/slayer/sql/dialects/__init__.py index e70aa437..8f946b54 100644 --- a/slayer/sql/dialects/__init__.py +++ b/slayer/sql/dialects/__init__.py @@ -14,8 +14,6 @@ from __future__ import annotations -from typing import Optional - from slayer.sql.dialects._tier2 import ( DatabricksDialect, OracleDialect, @@ -96,7 +94,7 @@ def get_dialect(sqlglot_name: str) -> SqlDialect: return _BY_SQLGLOT_NAME[sqlglot_name] -def dialect_for_ds_type(ds_type: Optional[str]) -> SqlDialect: +def dialect_for_ds_type(ds_type: str | None) -> SqlDialect: """Lenient lookup by datasource-config ``type`` string. Falls back to ``PostgresDialect`` for ``None``, empty, OR unknown diff --git a/slayer/sql/dialects/_tier2.py b/slayer/sql/dialects/_tier2.py index 686baea3..244b3ca8 100644 --- a/slayer/sql/dialects/_tier2.py +++ b/slayer/sql/dialects/_tier2.py @@ -20,7 +20,6 @@ from __future__ import annotations -from typing import Optional from slayer.sql.dialects.base import SqlDialect @@ -28,7 +27,7 @@ class RedshiftDialect(SqlDialect): sqlglot_name: str = "redshift" ds_type_aliases: frozenset[str] = frozenset({"redshift"}) - explain_prefix: Optional[str] = "EXPLAIN" + explain_prefix: str | None = "EXPLAIN" explain_postfix: str = "" log10_native: bool = True log2_native: bool = False @@ -37,7 +36,7 @@ class RedshiftDialect(SqlDialect): class TrinoDialect(SqlDialect): sqlglot_name: str = "trino" ds_type_aliases: frozenset[str] = frozenset({"trino"}) - explain_prefix: Optional[str] = "EXPLAIN ANALYZE" + explain_prefix: str | None = "EXPLAIN ANALYZE" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True @@ -47,7 +46,7 @@ class PrestoDialect(SqlDialect): sqlglot_name: str = "presto" # Athena uses the Presto dialect via this alias. ds_type_aliases: frozenset[str] = frozenset({"presto", "athena"}) - explain_prefix: Optional[str] = "EXPLAIN ANALYZE" + explain_prefix: str | None = "EXPLAIN ANALYZE" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True @@ -56,7 +55,7 @@ class PrestoDialect(SqlDialect): class DatabricksDialect(SqlDialect): sqlglot_name: str = "databricks" ds_type_aliases: frozenset[str] = frozenset({"databricks"}) - explain_prefix: Optional[str] = "EXPLAIN EXTENDED" + explain_prefix: str | None = "EXPLAIN EXTENDED" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True @@ -65,7 +64,7 @@ class DatabricksDialect(SqlDialect): class SparkDialect(SqlDialect): sqlglot_name: str = "spark" ds_type_aliases: frozenset[str] = frozenset({"spark"}) - explain_prefix: Optional[str] = "EXPLAIN EXTENDED" + explain_prefix: str | None = "EXPLAIN EXTENDED" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True @@ -74,7 +73,7 @@ class SparkDialect(SqlDialect): class OracleDialect(SqlDialect): sqlglot_name: str = "oracle" ds_type_aliases: frozenset[str] = frozenset({"oracle"}) - explain_prefix: Optional[str] = "EXPLAIN PLAN FOR" + explain_prefix: str | None = "EXPLAIN PLAN FOR" explain_postfix: str = "" # Oracle has neither LOG10 nor LOG2 as single-arg functions — keep # the canonical 2-arg LOG(base, x) form. diff --git a/slayer/sql/dialects/base.py b/slayer/sql/dialects/base.py index 596ec5be..a33da6cd 100644 --- a/slayer/sql/dialects/base.py +++ b/slayer/sql/dialects/base.py @@ -12,7 +12,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import TYPE_CHECKING, Any +from collections.abc import Callable from pydantic import BaseModel, ConfigDict from sqlglot import exp @@ -153,7 +154,7 @@ class SqlDialect(BaseModel): sqlglot_name: str = "postgres" ds_type_aliases: frozenset[str] = frozenset() - explain_prefix: Optional[str] = "EXPLAIN" + explain_prefix: str | None = "EXPLAIN" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True @@ -367,10 +368,10 @@ def emit_outer_wrap( *, inner_sql: str, public: list[str], - order: Optional[exp.Expression], - limit: Optional[exp.Expression], - offset_arg: Optional[exp.Expression], - parse: Optional[Callable[[str], exp.Expression]] = None, + order: exp.Expression | None, + limit: exp.Expression | None, + offset_arg: exp.Expression | None, + parse: Callable[[str], exp.Expression] | None = None, ) -> str: """Emit the DEV-1444 outer-projection wrap around ``inner_sql``. @@ -507,7 +508,7 @@ def build_explain_sql(self, sql: str) -> str: def build_connection_url( self, datasource: "DatasourceConfig", - ) -> Optional[str]: + ) -> str | None: """Hook: dialect-specific connection-string builder. Returning ``None`` (the default) means: defer to @@ -524,7 +525,7 @@ def build_engine( datasource: "DatasourceConfig", *, connection_string: str, - ) -> Optional["sa.Engine"]: + ) -> "sa.Engine | None": """Hook: build a dialect-specific SQLAlchemy engine. Returning ``None`` (the default) means: ``engine_factory`` falls @@ -549,7 +550,7 @@ def apply_session_overrides( """ return None - def statement_timeout_sql(self, timeout_seconds: int) -> Optional[str]: + def statement_timeout_sql(self, timeout_seconds: int) -> str | None: """Hook: SQL to set a per-statement timeout, or ``None`` if the dialect doesn't expose one or the existing client.py path handles it via a hardcoded branch (mysql / clickhouse / postgres). @@ -559,7 +560,7 @@ def statement_timeout_sql(self, timeout_seconds: int) -> Optional[str]: """ return None - def map_cursor_type_code(self, type_code: int) -> Optional[str]: + def map_cursor_type_code(self, type_code: int) -> str | None: """Hook: dialect-specific cursor-type-code → SLayer category (one of ``"number"``, ``"string"``, ``"time"``, ``"boolean"``). diff --git a/slayer/sql/dialects/bigquery.py b/slayer/sql/dialects/bigquery.py index ff86ba5a..b3437a06 100644 --- a/slayer/sql/dialects/bigquery.py +++ b/slayer/sql/dialects/bigquery.py @@ -27,7 +27,8 @@ import json import re -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import TYPE_CHECKING, Any +from collections.abc import Callable import sqlalchemy as sa from sqlglot import exp @@ -83,7 +84,7 @@ class BigqueryDialect(SqlDialect): sqlglot_name: str = "bigquery" ds_type_aliases: frozenset[str] = frozenset({"bigquery"}) # BigQuery has no SQL-level EXPLAIN. - explain_prefix: Optional[str] = None + explain_prefix: str | None = None explain_postfix: str = "" log10_native: bool = True log2_native: bool = True @@ -147,7 +148,7 @@ def build_engine( datasource: "DatasourceConfig", *, connection_string: str, - ) -> Optional["sa.Engine"]: + ) -> "sa.Engine | None": """Construct the SQLAlchemy engine with inline service-account JSON when ``DatasourceConfig.credentials_json`` is set. diff --git a/slayer/sql/dialects/clickhouse.py b/slayer/sql/dialects/clickhouse.py index b6ee744d..ce8f9f1d 100644 --- a/slayer/sql/dialects/clickhouse.py +++ b/slayer/sql/dialects/clickhouse.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import Callable, Optional +from collections.abc import Callable from sqlglot import exp @@ -17,7 +17,7 @@ class ClickhouseDialect(SqlDialect): sqlglot_name: str = "clickhouse" ds_type_aliases: frozenset[str] = frozenset({"clickhouse"}) - explain_prefix: Optional[str] = "EXPLAIN" + explain_prefix: str | None = "EXPLAIN" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True diff --git a/slayer/sql/dialects/duckdb.py b/slayer/sql/dialects/duckdb.py index 9831f62b..9e31c70c 100644 --- a/slayer/sql/dialects/duckdb.py +++ b/slayer/sql/dialects/duckdb.py @@ -7,7 +7,6 @@ from __future__ import annotations -from typing import Optional from slayer.sql.dialects.base import SqlDialect @@ -15,7 +14,7 @@ class DuckdbDialect(SqlDialect): sqlglot_name: str = "duckdb" ds_type_aliases: frozenset[str] = frozenset({"duckdb"}) - explain_prefix: Optional[str] = "EXPLAIN ANALYZE" + explain_prefix: str | None = "EXPLAIN ANALYZE" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True diff --git a/slayer/sql/dialects/mysql.py b/slayer/sql/dialects/mysql.py index 3b890bae..c9fd6974 100644 --- a/slayer/sql/dialects/mysql.py +++ b/slayer/sql/dialects/mysql.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Callable, Optional +from collections.abc import Callable from sqlglot import exp @@ -21,7 +21,7 @@ class MysqlDialect(SqlDialect): sqlglot_name: str = "mysql" ds_type_aliases: frozenset[str] = frozenset({"mysql", "mariadb"}) - explain_prefix: Optional[str] = "EXPLAIN FORMAT=JSON" + explain_prefix: str | None = "EXPLAIN FORMAT=JSON" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True diff --git a/slayer/sql/dialects/postgres.py b/slayer/sql/dialects/postgres.py index bf851afd..65909149 100644 --- a/slayer/sql/dialects/postgres.py +++ b/slayer/sql/dialects/postgres.py @@ -6,7 +6,6 @@ from __future__ import annotations -from typing import Optional from sqlglot import exp @@ -42,7 +41,7 @@ def _cast_round_arg_to_numeric(node: exp.Expression) -> exp.Expression: class PostgresDialect(SqlDialect): sqlglot_name: str = "postgres" ds_type_aliases: frozenset[str] = frozenset({"postgres", "postgresql"}) - explain_prefix: Optional[str] = "EXPLAIN ANALYZE" + explain_prefix: str | None = "EXPLAIN ANALYZE" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True diff --git a/slayer/sql/dialects/snowflake.py b/slayer/sql/dialects/snowflake.py index c1456206..c32cff2e 100644 --- a/slayer/sql/dialects/snowflake.py +++ b/slayer/sql/dialects/snowflake.py @@ -22,7 +22,7 @@ from __future__ import annotations import re as _re -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any from urllib.parse import quote import sqlalchemy as sa @@ -186,7 +186,7 @@ class so ``engine_factory`` / ``client`` stay dialect-agnostic. sqlglot_name: str = "snowflake" ds_type_aliases: frozenset[str] = frozenset({"snowflake"}) - explain_prefix: Optional[str] = "EXPLAIN USING JSON" + explain_prefix: str | None = "EXPLAIN USING JSON" explain_postfix: str = "" log10_native: bool = True # No native LOG2 — falls through to canonical ``LOG(2, x)`` form. @@ -199,7 +199,7 @@ class so ``engine_factory`` / ``client`` stay dialect-agnostic. def build_connection_url( self, datasource: "DatasourceConfig", - ) -> Optional[str]: + ) -> str | None: """Emit the sentinel URL when ``connection_name`` is set, otherwise build the full snowflake-sqlalchemy URL from inline fields. @@ -239,7 +239,7 @@ def build_engine( datasource: "DatasourceConfig", *, connection_string: str, - ) -> Optional["sa.Engine"]: + ) -> "sa.Engine | None": """When the sentinel URL is in play, route through ``creator=`` so ``snowflake.connector.connect(connection_name=...)`` drives the auth path. Otherwise return None to let ``engine_factory`` @@ -351,7 +351,7 @@ def apply_session_overrides( # Runtime statement hooks # ------------------------------------------------------------------ - def statement_timeout_sql(self, timeout_seconds: int) -> Optional[str]: + def statement_timeout_sql(self, timeout_seconds: int) -> str | None: """``ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = N``. Per-session setting; takes effect for every subsequent statement @@ -360,7 +360,7 @@ def statement_timeout_sql(self, timeout_seconds: int) -> Optional[str]: """ return f"ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = {timeout_seconds}" - def map_cursor_type_code(self, type_code: int) -> Optional[str]: + def map_cursor_type_code(self, type_code: int) -> str | None: """Map a snowflake-connector ``FieldType`` integer code to a SLayer category. Returns ``None`` for unknown codes so the caller can fall through to a default rather than mis-classify.""" diff --git a/slayer/sql/dialects/sqlite.py b/slayer/sql/dialects/sqlite.py index 5763feca..13f2994d 100644 --- a/slayer/sql/dialects/sqlite.py +++ b/slayer/sql/dialects/sqlite.py @@ -15,7 +15,7 @@ from __future__ import annotations import math -from typing import Callable, Optional +from collections.abc import Callable from sqlglot import exp @@ -91,7 +91,7 @@ def step(self, value) -> None: if value is not None: self._vals.append(value) - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if not self._vals: return None s = sorted(self._vals) @@ -107,7 +107,7 @@ class _PercentileContAgg: def __init__(self) -> None: self._vals: list[float] = [] - self._p: Optional[float] = None + self._p: float | None = None def step(self, value, p) -> None: if p is not None: @@ -118,7 +118,7 @@ def step(self, value, p) -> None: if value is not None: self._vals.append(value) - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if not self._vals or self._p is None: return None s = sorted(self._vals) @@ -136,7 +136,7 @@ class _PercentileDiscAgg: def __init__(self) -> None: self._vals: list[float] = [] - self._p: Optional[float] = None + self._p: float | None = None def step(self, value, p) -> None: if p is not None: @@ -191,7 +191,7 @@ def step(self, value) -> None: class _StddevSampAgg(_OneVarWelford): """Sample standard deviation. NULL when N <= 1.""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n <= 1: return None return math.sqrt(self._m2 / (self._n - 1)) @@ -200,7 +200,7 @@ def finalize(self) -> Optional[float]: class _StddevPopAgg(_OneVarWelford): """Population standard deviation. NULL at N=0; 0 at N=1.""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n == 0: return None if self._n == 1: @@ -211,7 +211,7 @@ def finalize(self) -> Optional[float]: class _VarSampAgg(_OneVarWelford): """Sample variance. NULL when N <= 1.""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n <= 1: return None return self._m2 / (self._n - 1) @@ -220,7 +220,7 @@ def finalize(self) -> Optional[float]: class _VarPopAgg(_OneVarWelford): """Population variance. NULL at N=0; 0 at N=1.""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n == 0: return None if self._n == 1: @@ -256,7 +256,7 @@ class _CorrAgg(_PairAgg): """Pearson correlation. NULL when fewer than 2 non-null pairs OR when either side has zero variance (matches Postgres CORR).""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n < 2: return None if self._m2_x == 0 or self._m2_y == 0: @@ -267,7 +267,7 @@ def finalize(self) -> Optional[float]: class _CovarSampAgg(_PairAgg): """Sample covariance. NULL when N <= 1.""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n <= 1: return None return self._c / (self._n - 1) @@ -276,7 +276,7 @@ def finalize(self) -> Optional[float]: class _CovarPopAgg(_PairAgg): """Population covariance. NULL at N=0; 0 at N=1.""" - def finalize(self) -> Optional[float]: + def finalize(self) -> float | None: if self._n == 0: return None if self._n == 1: @@ -402,7 +402,7 @@ def register_sqlite_udfs(dbapi_connection) -> None: class SqliteDialect(SqlDialect): sqlglot_name: str = "sqlite" ds_type_aliases: frozenset[str] = frozenset({"sqlite"}) - explain_prefix: Optional[str] = "EXPLAIN QUERY PLAN" + explain_prefix: str | None = "EXPLAIN QUERY PLAN" explain_postfix: str = "" log10_native: bool = True log2_native: bool = True diff --git a/slayer/sql/dialects/tsql.py b/slayer/sql/dialects/tsql.py index 72fa0483..a0543a76 100644 --- a/slayer/sql/dialects/tsql.py +++ b/slayer/sql/dialects/tsql.py @@ -27,7 +27,8 @@ from __future__ import annotations import re -from typing import Any, Callable, Optional +from typing import Any +from collections.abc import Callable import sqlglot from sqlglot import exp @@ -65,7 +66,7 @@ class TsqlDialect(SqlDialect): sqlglot_name: str = "tsql" ds_type_aliases: frozenset[str] = frozenset({"mssql", "sqlserver", "tsql"}) - explain_prefix: Optional[str] = "SET SHOWPLAN_ALL ON;" + explain_prefix: str | None = "SET SHOWPLAN_ALL ON;" explain_postfix: str = "; SET SHOWPLAN_ALL OFF" log10_native: bool = True log2_native: bool = False @@ -222,10 +223,10 @@ def emit_outer_wrap( *, inner_sql: str, public: list[str], - order: Optional[exp.Expression], - limit: Optional[exp.Expression], - offset_arg: Optional[exp.Expression], - parse: Optional[Callable[[str], exp.Expression]] = None, + order: exp.Expression | None, + limit: exp.Expression | None, + offset_arg: exp.Expression | None, + parse: Callable[[str], exp.Expression] | None = None, ) -> str: """T-SQL: hoist inner top-level CTEs to the outer statement AND transpose detached pagination to ``TOP`` / ``FETCH NEXT N ROWS diff --git a/slayer/sql/engine_factory.py b/slayer/sql/engine_factory.py index da1d2632..5cd46904 100644 --- a/slayer/sql/engine_factory.py +++ b/slayer/sql/engine_factory.py @@ -22,7 +22,6 @@ from __future__ import annotations import logging -from typing import Dict, Tuple import sqlalchemy as sa import sqlalchemy.event as sa_event @@ -35,7 +34,7 @@ # Engine cache. Key = (connection_string, runtime_fingerprint). -_engine_cache: Dict[Tuple[str, str], sa.Engine] = {} +_engine_cache: dict[tuple[str, str], sa.Engine] = {} def _runtime_fingerprint(datasource: DatasourceConfig) -> str: diff --git a/slayer/sql/generator.py b/slayer/sql/generator.py index c3d058ca..5cc0ddb4 100644 --- a/slayer/sql/generator.py +++ b/slayer/sql/generator.py @@ -8,7 +8,7 @@ import copy import logging import re -from typing import Any, Dict, List, Optional, Union +from typing import Any import sqlglot from sqlglot import exp @@ -23,7 +23,7 @@ from slayer.sql.dialects import SqlDialect, get_dialect -def _wrap_cast_for_type(expr: exp.Expression, dt: Optional[DataType]) -> exp.Expression: +def _wrap_cast_for_type(expr: exp.Expression, dt: DataType | None) -> exp.Expression: """DEV-1361: wrap ``expr`` in ``CAST(expr AS )`` so the declared SLayer ``DataType`` is enforced in emitted SQL. @@ -139,7 +139,7 @@ def _wrap_cast_for_type(expr: exp.Expression, dt: Optional[DataType]) -> exp.Exp ) -def _wrap_filter(sql_str: str, filter_sql: Optional[str]) -> str: +def _wrap_filter(sql_str: str, filter_sql: str | None) -> str: """Wrap ``sql_str`` in ``CASE WHEN filter_sql THEN ... END`` if a row-level filter is set; otherwise pass through unchanged. Used by the dialect-aware aggregate builders (``_build_percentile``, ``_build_stat_agg``, @@ -379,7 +379,7 @@ def _strip_trailing_pagination(sql: str) -> str: class SQLGenerator: """Generates SQL from an EnrichedQuery.""" - def __init__(self, dialect: Union[str, SqlDialect] = "postgres"): + def __init__(self, dialect: str | SqlDialect = "postgres"): if isinstance(dialect, SqlDialect): self._dialect: SqlDialect = dialect else: @@ -408,7 +408,7 @@ def _q(self, name: str) -> str: """ return exp.Identifier(this=name, quoted=True).sql(dialect=self.dialect) - def _parse(self, sql: str, *, dialect: Optional[str] = None) -> exp.Expression: + def _parse(self, sql: str, *, dialect: str | None = None) -> exp.Expression: """Parse ``sql`` via sqlglot, applying SLayer-specific AST rewrites. On SQLite, rewrites ``exp.JSONExtract`` to the function-call form so @@ -441,7 +441,7 @@ def _parse(self, sql: str, *, dialect: Optional[str] = None) -> exp.Expression: # SQLite / DuckDB output. return self._dialect.rewrite_target_ast(tree) - def _parse_predicate(self, sql: str, *, dialect: Optional[str] = None) -> exp.Expression: + def _parse_predicate(self, sql: str, *, dialect: str | None = None) -> exp.Expression: """Parse a bare WHERE/HAVING predicate expression (DEV-1378). ``sqlglot.parse_one(sql, dialect=...)`` falls back to a ``Command`` @@ -603,7 +603,7 @@ def _build_outer_wrap( self, *, inner_sql: str, - public: List[str], + public: list[str], order, limit, offset_arg, @@ -1856,7 +1856,7 @@ def _apply_order_limit(self, select: exp.Select, enriched: EnrichedQuery) -> exp col_name = self._resolve_order_column(col=col, enriched=enriched) order_col = exp.Column(this=exp.to_identifier(col_name, quoted=True)) ascending = order_item.direction == "asc" - ordered_kwargs: Dict[str, Any] = {"this": order_col, "desc": not ascending} + ordered_kwargs: dict[str, Any] = {"this": order_col, "desc": not ascending} if suppress_nulls_emulation: # T-SQL native: NULLS FIRST on ASC, NULLS LAST on DESC. ordered_kwargs["nulls_first"] = ascending @@ -2110,10 +2110,10 @@ def _rewrite_log_aliases(self, node: exp.Expression) -> exp.Expression: def _resolve_sql( self, - sql: Optional[str], + sql: str | None, name: str, model_name: str, - type: Optional[DataType] = None, + type: DataType | None = None, ) -> exp.Expression: """Resolve an enriched SQL expression to a sqlglot AST node. @@ -2162,7 +2162,7 @@ def _resolve_agg_param( ``_build_stat_agg`` (``other=``); mirrors ``weighted_avg``'s ``weight=`` flow. """ - raw: Optional[str] = None + raw: str | None = None if name in measure.agg_kwargs: raw = measure.agg_kwargs[name] _validate_agg_param_value(raw, name, agg_name) @@ -2184,10 +2184,10 @@ def _resolve_agg_param( def _build_agg( self, measure: EnrichedMeasure, - rn_suffix_map: Optional[dict[str, str]] = None, - default_time_col: Optional[str] = None, - filtered_rn_map: Optional[dict[str, str]] = None, - filtered_match_map: Optional[dict[str, str]] = None, + rn_suffix_map: dict[str, str] | None = None, + default_time_col: str | None = None, + filtered_rn_map: dict[str, str] | None = None, + filtered_match_map: dict[str, str] | None = None, ) -> tuple[exp.Expression, bool]: """Build an aggregation expression from an enriched measure.""" agg_name = measure.aggregation @@ -2444,7 +2444,7 @@ def _build_stat_agg(self, measure: "EnrichedMeasure") -> exp.Expression: # missing-required-param error takes priority over any dialect-specific # error when both conditions hold — the missing-param message points at # the actual user mistake. Closes Codex #5 on PR #82. - other_expr: Optional[str] = None + other_expr: str | None = None if agg_name in _TWO_ARG_STAT_AGGS: other_expr = _wrap_filter( self._resolve_agg_param(measure, name="other", agg_name=agg_name), @@ -2472,9 +2472,9 @@ def _build_stat_agg(self, measure: "EnrichedMeasure") -> exp.Expression: def _build_where_and_having( self, enriched: EnrichedQuery, - rn_suffix_map: Optional[dict[str, str]] = None, - filtered_rn_map: Optional[dict[str, str]] = None, - ) -> tuple[Optional[exp.Expression], Optional[exp.Expression]]: + rn_suffix_map: dict[str, str] | None = None, + filtered_rn_map: dict[str, str] | None = None, + ) -> tuple[exp.Expression | None, exp.Expression | None]: """Build WHERE and HAVING clauses from parsed filters. ParsedFilter objects have pre-built SQL strings. Column names are diff --git a/slayer/sql/sql_predicate.py b/slayer/sql/sql_predicate.py index a2d0bfdc..e2182919 100644 --- a/slayer/sql/sql_predicate.py +++ b/slayer/sql/sql_predicate.py @@ -19,7 +19,6 @@ from __future__ import annotations import re -from typing import List from slayer.core.formula import ALL_TRANSFORMS, ParsedFilter from slayer.core.refs import AGG_REF_RE @@ -60,7 +59,7 @@ def _reject_dsl_constructs(formula: str) -> None: ) -def _bare_column_refs(formula: str) -> List[str]: +def _bare_column_refs(formula: str) -> list[str]: """Best-effort extraction of column-like identifiers from a SQL-mode predicate, used so downstream join-detection sees the same shape it saw before DEV-1369 round 2 dropped sqlglot from this path. Skips @@ -69,7 +68,7 @@ def _bare_column_refs(formula: str) -> List[str]: ``status`` → ``status``) and excludes function-call heads. """ stripped = _strip_string_literals(formula) - refs: List[str] = [] + refs: list[str] = [] for match in re.finditer(r"\b[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*)*\b", stripped): token = match.group(0) # Skip if immediately followed by `(` (function call) diff --git a/slayer/sql/sqlite_introspect.py b/slayer/sql/sqlite_introspect.py index cc42fe10..2b71c3ef 100644 --- a/slayer/sql/sqlite_introspect.py +++ b/slayer/sql/sqlite_introspect.py @@ -26,7 +26,6 @@ import logging import math -from typing import Optional import sqlalchemy as sa from pydantic import BaseModel @@ -53,7 +52,7 @@ def _quote_sqlite_ident(name: str) -> str: return exp.to_identifier(name, quoted=True).sql(dialect="sqlite") -def _qualified_table(table: str, schema: Optional[str]) -> str: +def _qualified_table(table: str, schema: str | None) -> str: """Build a quoted, optionally schema-qualified SQLite table reference.""" tname = _quote_sqlite_ident(table) if schema: @@ -74,7 +73,7 @@ class _ProbeCounts(BaseModel): def _run_main_probe( *, conn: sa.engine.Connection, qtable: str, qcol: str, table: str, column: str, -) -> Optional[_ProbeCounts]: +) -> _ProbeCounts | None: """Run the main probe aggregate. Returns ``None`` on driver failure (after logging a WARNING).""" scan_cap_plus_one = PROBE_SCAN_CAP + 1 @@ -110,7 +109,7 @@ def _run_main_probe( def _classify_text_branch( *, conn: sa.engine.Connection, qtable: str, qcol: str, table: str, column: str, -) -> Optional[DataType]: +) -> DataType | None: """The ``n_text > 0`` decision branch — runs the bounded distinct-text coerce probe and classifies. Returns the verdict (DOUBLE / TEXT) or ``None`` on coerce-probe driver failure (after logging a WARNING). @@ -154,7 +153,7 @@ def _classify_text_branch( def _classify_no_evidence( *, counts: _ProbeCounts, table: str, column: str, -) -> Optional[DataType]: +) -> DataType | None: """The all-zero-evidence branch — saturated, empty, or all-NULL all decline to certify INT. Returns ``None`` after logging a WARNING for saturated samples; returns ``None`` silently for empty / all-NULL; @@ -179,8 +178,8 @@ def probe_sqlite_integer_column( conn: sa.engine.Connection, table: str, column: str, - schema: Optional[str] = None, -) -> Optional[DataType]: + schema: str | None = None, +) -> DataType | None: """Probe one SQLite column's actual storage classes and return the appropriate :class:`DataType`. diff --git a/slayer/storage/base.py b/slayer/storage/base.py index 9d6521c2..c7cd16bc 100644 --- a/slayer/storage/base.py +++ b/slayer/storage/base.py @@ -5,7 +5,8 @@ import sys from abc import ABC, abstractmethod from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any +from collections.abc import Callable from slayer.core.errors import AmbiguousModelError, MemoryNotFoundError from slayer.core.models import DatasourceConfig, SlayerModel @@ -25,11 +26,11 @@ def _write_sample_fields( - col: Dict[str, Any], + col: dict[str, Any], *, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: """Apply the DEV-1375 + DEV-1480 sample-field write convention to a column dict in place: ``None`` pops the corresponding key, non-None @@ -199,7 +200,7 @@ async def _save_model_impl(self, model: SlayerModel) -> None: """ @abstractmethod - async def _list_all_model_identities(self) -> List[Tuple[str, str]]: + async def _list_all_model_identities(self) -> list[tuple[str, str]]: """Return every saved ``(data_source, name)`` pair. Backends override this with whatever is cheapest (filesystem walk, @@ -211,13 +212,13 @@ async def _list_all_model_identities(self) -> List[Tuple[str, str]]: async def get_model( self, name: str, - data_source: Optional[str] = None, - ) -> Optional[SlayerModel]: ... + data_source: str | None = None, + ) -> SlayerModel | None: ... async def delete_model( self, name: str, - data_source: Optional[str] = None, + data_source: str | None = None, ) -> bool: """Delete one model by ``(data_source, name)`` and cascade-delete every embedding row tagged with that model's canonical prefix. @@ -259,9 +260,9 @@ async def update_column_sampled( data_source: str, model_name: str, column_name: str, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: """Patch a single column's sample-value fields in-place (DEV-1375 + DEV-1480). @@ -281,8 +282,8 @@ async def _resolve_target_or_none( self, name: str, *, - data_source: Optional[str], - ) -> Optional[Tuple[str, str]]: + data_source: str | None, + ) -> tuple[str, str] | None: """Sanitize inputs and resolve a bare ``name`` to its ``(data_source, name)`` identity via the priority list. @@ -395,10 +396,10 @@ async def _migrate_and_refine_on_load( async def save_datasource(self, datasource: DatasourceConfig) -> None: ... @abstractmethod - async def get_datasource(self, name: str) -> Optional[DatasourceConfig]: ... + async def get_datasource(self, name: str) -> DatasourceConfig | None: ... @abstractmethod - async def list_datasources(self) -> List[str]: ... + async def list_datasources(self) -> list[str]: ... async def delete_datasource(self, name: str) -> bool: """Delete the datasource config and cascade-delete every embedding @@ -434,7 +435,7 @@ async def _delete_datasource_row(self, name: str) -> bool: # ---- datasource priority (bare-name disambiguation) ------------------- @abstractmethod - async def get_datasource_priority(self) -> List[str]: + async def get_datasource_priority(self) -> list[str]: """Return the configured priority order (most-preferred first). Empty list = no priority configured; bare-name lookups raise @@ -442,11 +443,11 @@ async def get_datasource_priority(self) -> List[str]: """ @abstractmethod - async def _set_datasource_priority_raw(self, priority: List[str]) -> None: + async def _set_datasource_priority_raw(self, priority: list[str]) -> None: """Persist the priority list verbatim. Validation happens in the public ``set_datasource_priority`` wrapper below.""" - async def set_datasource_priority(self, priority: List[str]) -> None: + async def set_datasource_priority(self, priority: list[str]) -> None: """Validate and persist the datasource priority list. Each entry must already exist as a saved ``DatasourceConfig``; @@ -465,7 +466,7 @@ async def set_datasource_priority(self, priority: List[str]) -> None: # ---- list_models with auto-detect or required arg ---------------------- - async def list_models(self, data_source: Optional[str] = None) -> List[str]: + async def list_models(self, data_source: str | None = None) -> list[str]: """List model names within a single datasource. Resolution rules: @@ -511,8 +512,8 @@ async def resolve_model_identity( self, name: str, *, - prefer_data_source: Optional[str] = None, - ) -> Optional[Tuple[str, str]]: + prefer_data_source: str | None = None, + ) -> tuple[str, str] | None: """Resolve a bare model name to a ``(data_source, name)`` tuple. * No matches → ``None``. @@ -560,10 +561,10 @@ async def _save_memory_row(self, memory: Memory) -> None: be preserved when present.""" @abstractmethod - async def _get_memory_row(self, memory_id: str) -> Optional[Memory]: + async def _get_memory_row(self, memory_id: str) -> Memory | None: """Read a ``Memory`` by id; return ``None`` when not present.""" - async def get_memory_row(self, memory_id: str) -> Optional[Memory]: + async def get_memory_row(self, memory_id: str) -> Memory | None: """Non-raising existence check / fetch. Public so the resolver and the ingest-time cleanup pass can probe without catching ``MemoryNotFoundError``.""" @@ -571,8 +572,8 @@ async def get_memory_row(self, memory_id: str) -> Optional[Memory]: @abstractmethod async def _list_memories_rows( - self, *, entities: Optional[List[str]] - ) -> List[Memory]: + self, *, entities: list[str] | None + ) -> list[Memory]: """Return every ``Memory`` whose stored entity set has non-empty intersection with ``entities``. ``entities=None`` returns all rows. ``entities=[]`` returns ``[]`` (intersection with the empty set is @@ -594,10 +595,10 @@ async def save_memory( self, *, learning: str, - entities: List[str], - query: Optional[SlayerQuery] = None, - id: Optional[str] = None, # noqa: A002 — public kwarg matching MCP / REST - description: Optional[str] = None, + entities: list[str], + query: SlayerQuery | None = None, + id: str | None = None, # noqa: A002 — public kwarg matching MCP / REST + description: str | None = None, ) -> Memory: """Persist a memory. @@ -620,7 +621,7 @@ async def save_memory( else: assigned_id = await self._next_memory_seq() preserved_created_at = None - kwargs: Dict[str, Any] = { + kwargs: dict[str, Any] = { "id": assigned_id, "learning": learning, "description": description, @@ -640,8 +641,8 @@ async def get_memory(self, memory_id: str) -> Memory: return row async def list_memories( - self, *, entities: Optional[List[str]] = None - ) -> List[Memory]: + self, *, entities: list[str] | None = None + ) -> list[Memory]: return await self._list_memories_rows(entities=entities) async def delete_memory(self, memory_id: str) -> None: @@ -782,13 +783,13 @@ async def save_embedding(self, row: Embedding) -> None: @abstractmethod async def get_embedding( self, *, canonical_id: str, embedding_model_name: str, - ) -> Optional[Embedding]: + ) -> Embedding | None: """Fetch one embedding row; ``None`` when no row matches.""" @abstractmethod async def list_embeddings( self, *, embedding_model_name: str, - ) -> List[Embedding]: + ) -> list[Embedding]: """Return every row for ``embedding_model_name``. Used by the search service to load the entire corpus into a numpy matrix.""" @@ -815,7 +816,7 @@ async def delete_embeddings_for_canonical( # working unchanged; the bundled backends override these to issue one # round-trip via :class:`SidecarEmbeddingStore`. - async def save_embeddings(self, rows: List[Embedding]) -> None: + async def save_embeddings(self, rows: list[Embedding]) -> None: """Persist many embedding rows in one round-trip. Default implementation calls :meth:`save_embedding` for each row.""" for row in rows: @@ -824,14 +825,14 @@ async def save_embeddings(self, rows: List[Embedding]) -> None: async def get_embeddings_for_canonical_ids( self, *, - canonical_ids: List[str], + canonical_ids: list[str], embedding_model_name: str, - ) -> Dict[str, "Embedding"]: + ) -> dict[str, "Embedding"]: """Fetch every embedding row in ``canonical_ids`` under the given ``embedding_model_name`` in one round-trip. Returns a dict keyed by ``canonical_id``; missing ids are simply absent from the dict. Default implementation calls :meth:`get_embedding` for each id.""" - out: Dict[str, Embedding] = {} + out: dict[str, Embedding] = {} for canonical_id in canonical_ids: row = await self.get_embedding( canonical_id=canonical_id, @@ -846,7 +847,7 @@ async def get_embeddings_for_canonical_ids( # Storage factory with pluggable registry # --------------------------------------------------------------------------- -_STORAGE_REGISTRY: Dict[str, Callable[[str], StorageBackend]] = {} +_STORAGE_REGISTRY: dict[str, Callable[[str], StorageBackend]] = {} def register_storage(scheme: str, factory: Callable[[str], StorageBackend]) -> None: diff --git a/slayer/storage/join_sync.py b/slayer/storage/join_sync.py index 0fe82b18..ddfd76b5 100644 --- a/slayer/storage/join_sync.py +++ b/slayer/storage/join_sync.py @@ -12,7 +12,6 @@ ``data_source`` only — cross-datasource joins are never auto-mirrored. """ -from typing import Dict, List, Optional, Set, Tuple from slayer.core.enums import JoinType from slayer.core.models import DatasourceConfig, ModelJoin, SlayerModel @@ -60,7 +59,7 @@ async def _mirror_inner_joins(model: SlayerModel, storage: StorageBackend) -> No async def _remove_inner_joins_to( source_name: str, source_data_source: str, - target_names: Set[str], + target_names: set[str], storage: StorageBackend, ) -> None: """Remove inner joins that point back to *source_name* from each target. @@ -158,7 +157,7 @@ async def _save_model_impl(self, model: SlayerModel) -> None: async def delete_model( self, name: str, - data_source: Optional[str] = None, + data_source: str | None = None, ) -> bool: await self._ensure_reconciled() if data_source is None: @@ -200,15 +199,15 @@ async def _delete_datasource_row(self, name: str) -> bool: # -- pure delegation -------------------------------------------------- - async def _list_all_model_identities(self) -> List[Tuple[str, str]]: + async def _list_all_model_identities(self) -> list[tuple[str, str]]: await self._ensure_reconciled() return await self._inner._list_all_model_identities() async def get_model( self, name: str, - data_source: Optional[str] = None, - ) -> Optional[SlayerModel]: + data_source: str | None = None, + ) -> SlayerModel | None: await self._ensure_reconciled() return await self._inner.get_model(name, data_source=data_source) @@ -218,9 +217,9 @@ async def update_column_sampled( data_source: str, model_name: str, column_name: str, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: return await self._inner.update_column_sampled( data_source=data_source, @@ -234,19 +233,19 @@ async def update_column_sampled( async def save_datasource(self, datasource: DatasourceConfig) -> None: return await self._inner.save_datasource(datasource) - async def get_datasource(self, name: str) -> Optional[DatasourceConfig]: + async def get_datasource(self, name: str) -> DatasourceConfig | None: return await self._inner.get_datasource(name) - async def list_datasources(self) -> List[str]: + async def list_datasources(self) -> list[str]: return await self._inner.list_datasources() async def delete_datasource(self, name: str) -> bool: return await self._inner.delete_datasource(name) - async def get_datasource_priority(self) -> List[str]: + async def get_datasource_priority(self) -> list[str]: return await self._inner.get_datasource_priority() - async def _set_datasource_priority_raw(self, priority: List[str]) -> None: + async def _set_datasource_priority_raw(self, priority: list[str]) -> None: return await self._inner._set_datasource_priority_raw(priority) # -- memories (DEV-1357 v2 / DEV-1428) — pure delegation -------------- @@ -254,12 +253,12 @@ async def _set_datasource_priority_raw(self, priority: List[str]) -> None: async def _save_memory_row(self, memory: Memory) -> None: await self._inner._save_memory_row(memory) - async def _get_memory_row(self, memory_id: str) -> Optional[Memory]: + async def _get_memory_row(self, memory_id: str) -> Memory | None: return await self._inner._get_memory_row(memory_id) async def _list_memories_rows( - self, *, entities: Optional[List[str]] - ) -> List[Memory]: + self, *, entities: list[str] | None + ) -> list[Memory]: return await self._inner._list_memories_rows(entities=entities) async def _delete_memory_row(self, memory_id: str) -> bool: @@ -280,12 +279,12 @@ async def strip_dangling_entities_from_memories( async def save_embedding(self, row: Embedding) -> None: await self._inner.save_embedding(row) - async def save_embeddings(self, rows: List[Embedding]) -> None: + async def save_embeddings(self, rows: list[Embedding]) -> None: await self._inner.save_embeddings(rows) async def get_embedding( self, *, canonical_id: str, embedding_model_name: str, - ) -> Optional[Embedding]: + ) -> Embedding | None: return await self._inner.get_embedding( canonical_id=canonical_id, embedding_model_name=embedding_model_name, @@ -294,9 +293,9 @@ async def get_embedding( async def get_embeddings_for_canonical_ids( self, *, - canonical_ids: List[str], + canonical_ids: list[str], embedding_model_name: str, - ) -> Dict[str, Embedding]: + ) -> dict[str, Embedding]: return await self._inner.get_embeddings_for_canonical_ids( canonical_ids=canonical_ids, embedding_model_name=embedding_model_name, @@ -304,7 +303,7 @@ async def get_embeddings_for_canonical_ids( async def list_embeddings( self, *, embedding_model_name: str, - ) -> List[Embedding]: + ) -> list[Embedding]: return await self._inner.list_embeddings( embedding_model_name=embedding_model_name, ) diff --git a/slayer/storage/migrations.py b/slayer/storage/migrations.py index 6a179307..729ac02c 100644 --- a/slayer/storage/migrations.py +++ b/slayer/storage/migrations.py @@ -11,10 +11,11 @@ version that ``.save_*()`` will write today. """ -from typing import Any, Callable, Dict, Tuple +from typing import Any +from collections.abc import Callable # Per-entity current version. Bump independently when an entity's schema changes. -CURRENT_VERSIONS: Dict[str, int] = { +CURRENT_VERSIONS: dict[str, int] = { "SlayerModel": 7, "SlayerQuery": 3, "DatasourceConfig": 2, @@ -23,7 +24,7 @@ } # Registry: (entity_name, source_version) -> converter producing source_version+1. -_REGISTRY: Dict[Tuple[str, int], Callable[[dict], dict]] = {} +_REGISTRY: dict[tuple[str, int], Callable[[dict], dict]] = {} def register_migration( diff --git a/slayer/storage/sidecar_embedding_store.py b/slayer/storage/sidecar_embedding_store.py index 5cc468af..ff2297cb 100644 --- a/slayer/storage/sidecar_embedding_store.py +++ b/slayer/storage/sidecar_embedding_store.py @@ -23,7 +23,6 @@ import asyncio import json import sqlite3 -from typing import Dict, List, Optional, Tuple from slayer.embeddings.models import Embedding @@ -62,7 +61,7 @@ def _init_db(self) -> None: # ------------------------------------------------------------------ @staticmethod - def _row_tuple(row: Embedding) -> Tuple[str, str, str, str, str, str]: + def _row_tuple(row: Embedding) -> tuple[str, str, str, str, str, str]: return ( row.canonical_id, row.embedding_model_name, @@ -73,7 +72,7 @@ def _row_tuple(row: Embedding) -> Tuple[str, str, str, str, str, str]: ) @staticmethod - def _row_from_db(raw: Tuple[str, str, str, str, str, str]) -> Embedding: + def _row_from_db(raw: tuple[str, str, str, str, str, str]) -> Embedding: return Embedding.model_validate({ "canonical_id": raw[0], "embedding_model_name": raw[1], @@ -93,7 +92,7 @@ def _save_sync(self, row: Embedding) -> None: self._row_tuple(row), ) - def _save_many_sync(self, rows: List[Embedding]) -> None: + def _save_many_sync(self, rows: list[Embedding]) -> None: with sqlite3.connect(self.db_path) as conn: conn.executemany( "INSERT OR REPLACE INTO embeddings " @@ -105,7 +104,7 @@ def _save_many_sync(self, rows: List[Embedding]) -> None: def _get_sync( self, canonical_id: str, embedding_model_name: str, - ) -> Optional[Tuple[str, str, str, str, str, str]]: + ) -> tuple[str, str, str, str, str, str] | None: with sqlite3.connect(self.db_path) as conn: row = conn.execute( "SELECT canonical_id, embedding_model_name, entity_kind, " @@ -125,10 +124,10 @@ def _get_sync( def _get_many_sync( self, - canonical_ids: List[str], + canonical_ids: list[str], embedding_model_name: str, - ) -> List[Tuple[str, str, str, str, str, str]]: - rows: List[Tuple[str, str, str, str, str, str]] = [] + ) -> list[tuple[str, str, str, str, str, str]]: + rows: list[tuple[str, str, str, str, str, str]] = [] with sqlite3.connect(self.db_path) as conn: for start in range(0, len(canonical_ids), self._GET_MANY_CHUNK_SIZE): chunk = canonical_ids[start : start + self._GET_MANY_CHUNK_SIZE] @@ -144,7 +143,7 @@ def _get_many_sync( def _list_sync( self, embedding_model_name: str, - ) -> List[Tuple[str, str, str, str, str, str]]: + ) -> list[tuple[str, str, str, str, str, str]]: with sqlite3.connect(self.db_path) as conn: rows = conn.execute( "SELECT canonical_id, embedding_model_name, entity_kind, " @@ -181,14 +180,14 @@ def _delete_by_prefix_sync(self, prefix: str) -> int: async def save(self, row: Embedding) -> None: await asyncio.to_thread(self._save_sync, row) - async def save_many(self, rows: List[Embedding]) -> None: + async def save_many(self, rows: list[Embedding]) -> None: if not rows: return await asyncio.to_thread(self._save_many_sync, list(rows)) async def get( self, *, canonical_id: str, embedding_model_name: str, - ) -> Optional[Embedding]: + ) -> Embedding | None: raw = await asyncio.to_thread( self._get_sync, canonical_id, embedding_model_name, ) @@ -199,9 +198,9 @@ async def get( async def get_many( self, *, - canonical_ids: List[str], + canonical_ids: list[str], embedding_model_name: str, - ) -> Dict[str, Embedding]: + ) -> dict[str, Embedding]: if not canonical_ids: return {} raws = await asyncio.to_thread( @@ -211,7 +210,7 @@ async def get_many( async def list_for_model( self, *, embedding_model_name: str, - ) -> List[Embedding]: + ) -> list[Embedding]: raws = await asyncio.to_thread( self._list_sync, embedding_model_name, ) @@ -243,12 +242,12 @@ class SidecarEmbeddingsMixin: async def save_embedding(self, row: Embedding) -> None: await self._embeddings_store.save(row) - async def save_embeddings(self, rows: List[Embedding]) -> None: + async def save_embeddings(self, rows: list[Embedding]) -> None: await self._embeddings_store.save_many(list(rows)) async def get_embedding( self, *, canonical_id: str, embedding_model_name: str, - ) -> Optional[Embedding]: + ) -> Embedding | None: return await self._embeddings_store.get( canonical_id=canonical_id, embedding_model_name=embedding_model_name, @@ -257,9 +256,9 @@ async def get_embedding( async def get_embeddings_for_canonical_ids( self, *, - canonical_ids: List[str], + canonical_ids: list[str], embedding_model_name: str, - ) -> Dict[str, Embedding]: + ) -> dict[str, Embedding]: return await self._embeddings_store.get_many( canonical_ids=list(canonical_ids), embedding_model_name=embedding_model_name, @@ -267,7 +266,7 @@ async def get_embeddings_for_canonical_ids( async def list_embeddings( self, *, embedding_model_name: str, - ) -> List[Embedding]: + ) -> list[Embedding]: return await self._embeddings_store.list_for_model( embedding_model_name=embedding_model_name, ) diff --git a/slayer/storage/sqlite_storage.py b/slayer/storage/sqlite_storage.py index 60cf9c8f..534750df 100644 --- a/slayer/storage/sqlite_storage.py +++ b/slayer/storage/sqlite_storage.py @@ -11,7 +11,7 @@ import json import os import sqlite3 -from typing import Any, Dict, List, Optional, Tuple +from typing import Any from slayer.core.models import DatasourceConfig, SlayerModel from slayer.core.query import SlayerQuery @@ -64,7 +64,7 @@ def _migrate_memories_to_text_pk(self) -> None: # Combine into one transaction so the cross-table FK rebuild # is atomic. Each ALTER/INSERT/DROP gated by the per-table # flags above; an empty branch is a no-op (no DDL emitted). - script_parts: List[str] = ["BEGIN;"] + script_parts: list[str] = ["BEGIN;"] if memories_needs_rebuild: script_parts.append( "ALTER TABLE memories RENAME TO _memories_legacy;" @@ -189,14 +189,14 @@ def _save_model_sync(self, model: SlayerModel) -> None: (model.data_source, model.name, data), ) - def _list_all_identities_sync(self) -> List[Tuple[str, str]]: + def _list_all_identities_sync(self) -> list[tuple[str, str]]: with sqlite3.connect(self.db_path) as conn: rows = conn.execute( "SELECT data_source, name FROM models ORDER BY data_source, name" ).fetchall() return [(r[0], r[1]) for r in rows] - def _get_model_sync(self, data_source: str, name: str) -> Optional[str]: + def _get_model_sync(self, data_source: str, name: str) -> str | None: with sqlite3.connect(self.db_path) as conn: row = conn.execute( "SELECT data FROM models WHERE data_source = ? AND name = ?", @@ -220,14 +220,14 @@ def _save_datasource_sync(self, datasource: DatasourceConfig) -> None: (datasource.name, data), ) - def _get_datasource_sync(self, name: str) -> Optional[str]: + def _get_datasource_sync(self, name: str) -> str | None: with sqlite3.connect(self.db_path) as conn: row = conn.execute( "SELECT data FROM datasources WHERE name = ?", (name,) ).fetchone() return row[0] if row else None - def _list_datasources_sync(self) -> List[str]: + def _list_datasources_sync(self) -> list[str]: with sqlite3.connect(self.db_path) as conn: rows = conn.execute( "SELECT name FROM datasources ORDER BY name" @@ -241,7 +241,7 @@ def _delete_datasource_sync(self, name: str) -> bool: ) return cursor.rowcount > 0 - def _get_priority_sync(self) -> List[str]: + def _get_priority_sync(self) -> list[str]: with sqlite3.connect(self.db_path) as conn: row = conn.execute( "SELECT value FROM settings WHERE key = ?", (_PRIORITY_KEY,) @@ -256,7 +256,7 @@ def _get_priority_sync(self) -> List[str]: return [] return [str(p) for p in value] - def _set_priority_sync(self, priority: List[str]) -> None: + def _set_priority_sync(self, priority: list[str]) -> None: with sqlite3.connect(self.db_path) as conn: conn.execute( "INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)", @@ -268,14 +268,14 @@ def _set_priority_sync(self, priority: List[str]) -> None: async def _save_model_impl(self, model: SlayerModel) -> None: await asyncio.to_thread(self._save_model_sync, model) - async def _list_all_model_identities(self) -> List[Tuple[str, str]]: + async def _list_all_model_identities(self) -> list[tuple[str, str]]: return await asyncio.to_thread(self._list_all_identities_sync) async def get_model( self, name: str, - data_source: Optional[str] = None, - ) -> Optional[SlayerModel]: + data_source: str | None = None, + ) -> SlayerModel | None: target = await self._resolve_target_or_none(name, data_source=data_source) if target is None: return None @@ -295,9 +295,9 @@ async def _delete_model_row( def _update_column_sampled_sync( self, *, data_source: str, model_name: str, - column_name: str, sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + column_name: str, sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: with sqlite3.connect(self.db_path) as conn: row = conn.execute( @@ -336,9 +336,9 @@ async def update_column_sampled( data_source: str, model_name: str, column_name: str, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: await asyncio.to_thread( self._update_column_sampled_sync, @@ -350,7 +350,7 @@ async def update_column_sampled( async def save_datasource(self, datasource: DatasourceConfig) -> None: await asyncio.to_thread(self._save_datasource_sync, datasource) - async def get_datasource(self, name: str) -> Optional[DatasourceConfig]: + async def get_datasource(self, name: str) -> DatasourceConfig | None: # DEV-1405: sanitize the raw name. Mirrors the YAML backend; the # SQLite lookup is parameterised so injection isn't the risk — # validation here keeps the public ABC contract uniform across @@ -362,16 +362,16 @@ async def get_datasource(self, name: str) -> Optional[DatasourceConfig]: ds = DatasourceConfig.model_validate(json.loads(raw)) return ds.resolve_env_vars() - async def list_datasources(self) -> List[str]: + async def list_datasources(self) -> list[str]: return await asyncio.to_thread(self._list_datasources_sync) async def _delete_datasource_row(self, name: str) -> bool: return await asyncio.to_thread(self._delete_datasource_sync, name) - async def get_datasource_priority(self) -> List[str]: + async def get_datasource_priority(self) -> list[str]: return await asyncio.to_thread(self._get_priority_sync) - async def _set_datasource_priority_raw(self, priority: List[str]) -> None: + async def _set_datasource_priority_raw(self, priority: list[str]) -> None: await asyncio.to_thread(self._set_priority_sync, list(priority)) # ---- memories (DEV-1357 v2) ------------------------------------------- @@ -402,11 +402,11 @@ def _is_int_shaped_id(value: Any) -> bool: def _save_memory_atomic_sync( self, *, - memory_id: Optional[str], + memory_id: str | None, learning: str, - entities: List[str], - query: Optional[SlayerQuery], - description: Optional[str] = None, + entities: list[str], + query: SlayerQuery | None, + description: str | None = None, ) -> Memory: """Reserve / accept an id and persist the new memory inside one SQLite transaction. Returns the persisted :class:`Memory`. @@ -443,7 +443,7 @@ def _save_memory_atomic_sync( json.loads(existing_row[0]) ) preserved_created_at = existing_memory.created_at - kwargs: Dict[str, Any] = { + kwargs: dict[str, Any] = { "id": memory_id, "learning": learning, "description": description, @@ -492,10 +492,10 @@ async def save_memory( self, *, learning: str, - entities: List[str], - query: Optional[SlayerQuery] = None, - id: Optional[str] = None, # noqa: A002 — public kwarg - description: Optional[str] = None, + entities: list[str], + query: SlayerQuery | None = None, + id: str | None = None, # noqa: A002 — public kwarg + description: str | None = None, ) -> Memory: if id is not None: _validate_memory_id_charset(id) @@ -543,20 +543,20 @@ def _next_memory_seq_sync(self) -> str: async def _next_memory_seq(self) -> str: return await asyncio.to_thread(self._next_memory_seq_sync) - def _get_memory_sync(self, memory_id: str) -> Optional[str]: + def _get_memory_sync(self, memory_id: str) -> str | None: with sqlite3.connect(self.db_path) as conn: row = conn.execute( "SELECT data FROM memories WHERE id = ?", (memory_id,) ).fetchone() return row[0] if row else None - async def _get_memory_row(self, memory_id: str) -> Optional[Memory]: + async def _get_memory_row(self, memory_id: str) -> Memory | None: raw = await asyncio.to_thread(self._get_memory_sync, memory_id) return Memory.model_validate(json.loads(raw)) if raw else None def _list_memories_sync( - self, entities: Optional[List[str]] - ) -> List[str]: + self, entities: list[str] | None + ) -> list[str]: with sqlite3.connect(self.db_path) as conn: if entities is None: rows = conn.execute( @@ -576,8 +576,8 @@ def _list_memories_sync( return [r[0] for r in rows] async def _list_memories_rows( - self, *, entities: Optional[List[str]] - ) -> List[Memory]: + self, *, entities: list[str] | None + ) -> list[Memory]: raws = await asyncio.to_thread(self._list_memories_sync, entities) return [Memory.model_validate(json.loads(r)) for r in raws] diff --git a/slayer/storage/type_refinement.py b/slayer/storage/type_refinement.py index 58ba5194..eae213da 100644 --- a/slayer/storage/type_refinement.py +++ b/slayer/storage/type_refinement.py @@ -26,7 +26,6 @@ from __future__ import annotations import logging -from typing import Optional, Tuple from slayer.core.enums import DataType @@ -36,7 +35,7 @@ logger = logging.getLogger(__name__) -def _column_is_base(sql: Optional[str]) -> bool: +def _column_is_base(sql: str | None) -> bool: """A column whose ``sql`` is ``None`` or a single bare identifier is a "base" column — it claims a live database column. Derived expressions (``amount * 2``, ``length(name)``, ``customers.region``, etc.) are not. @@ -54,7 +53,7 @@ def _is_sqlite_datasource(datasource: DatasourceConfig) -> bool: return (datasource.type or "").lower() == "sqlite" -def _is_auto_default_integer_format_dict(fmt: Optional[dict]) -> bool: +def _is_auto_default_integer_format_dict(fmt: dict | None) -> bool: """Dict-shape mirror of :func:`slayer.engine.ingestion._is_auto_default_integer_format` — detects the auto-ingested ``NumberFormat(type=INTEGER)`` default so @@ -71,7 +70,7 @@ def _is_auto_default_integer_format_dict(fmt: Optional[dict]) -> bool: ) -def _format_for_widened_type_dict(verdict: DataType) -> Optional[dict]: +def _format_for_widened_type_dict(verdict: DataType) -> dict | None: if verdict is DataType.DOUBLE: return {"type": "float"} return None # TEXT clears format @@ -135,7 +134,7 @@ def has_sqlite_widenable_columns(d: dict) -> bool: def _parse_sql_table_with_default_schema( sql_table: str, datasource: DatasourceConfig, -) -> Tuple[Optional[str], str]: +) -> tuple[str | None, str]: """Split ``sql_table`` into ``(schema, table)``, falling back to ``datasource.schema_name`` when the name is unqualified. This honours attached SQLite schemas instead of silently using ``main``. @@ -148,9 +147,9 @@ def _parse_sql_table_with_default_schema( def _safe_probe( - *, conn, table: str, column: str, schema: Optional[str], + *, conn, table: str, column: str, schema: str | None, sql_table: str, persisted_type: str, -) -> Optional[DataType]: +) -> DataType | None: """Run the probe with a defence-in-depth try/except so the refinement loop never aborts on an unexpected exception.""" from slayer.sql.sqlite_introspect import probe_sqlite_integer_column @@ -206,7 +205,7 @@ def _apply_format_flip( def _refine_one_column( - *, conn, col: dict, d: dict, table_name: str, schema_name: Optional[str], + *, conn, col: dict, d: dict, table_name: str, schema_name: str | None, sql_table: str, ) -> bool: """Run the probe for one column and apply the verdict. Returns True if diff --git a/slayer/storage/v2_datasource_migration.py b/slayer/storage/v2_datasource_migration.py index e1384eb4..be91aa94 100644 --- a/slayer/storage/v2_datasource_migration.py +++ b/slayer/storage/v2_datasource_migration.py @@ -11,12 +11,12 @@ Other dialects ignore the new fields entirely. """ -from typing import Any, Dict +from typing import Any from slayer.storage.migrations import register_migration @register_migration("DatasourceConfig", 1) -def _datasource_v1_to_v2(data: Dict[str, Any]) -> Dict[str, Any]: +def _datasource_v1_to_v2(data: dict[str, Any]) -> dict[str, Any]: """v1 → v2: no-op forward; new fields default to None via Pydantic.""" return data diff --git a/slayer/storage/v2_memory_migration.py b/slayer/storage/v2_memory_migration.py index 82f1047f..b7c4717b 100644 --- a/slayer/storage/v2_memory_migration.py +++ b/slayer/storage/v2_memory_migration.py @@ -13,13 +13,13 @@ this converter. """ -from typing import Any, Dict +from typing import Any from slayer.storage.migrations import register_migration @register_migration("Memory", 1) -def _memory_v1_to_v2(data: Dict[str, Any]) -> Dict[str, Any]: +def _memory_v1_to_v2(data: dict[str, Any]) -> dict[str, Any]: """v1 → v2: stringify the id field. Default sentinel preserved.""" raw_id = data.get("id", "") if isinstance(raw_id, bool): diff --git a/slayer/storage/v4_migration.py b/slayer/storage/v4_migration.py index 3bf800c2..7970a38e 100644 --- a/slayer/storage/v4_migration.py +++ b/slayer/storage/v4_migration.py @@ -28,7 +28,6 @@ import json import os import sqlite3 -from typing import List, Optional import yaml @@ -70,7 +69,7 @@ def _model_v3_to_v4(data: dict) -> dict: def _resolve_orphan_data_source( *, name: str, - available_datasources: List[str], + available_datasources: list[str], ) -> str: """Either auto-fill from the only datasource present, or raise.""" if len(available_datasources) == 1: @@ -89,7 +88,7 @@ def _resolve_orphan_data_source( # --------------------------------------------------------------------------- -def _yaml_list_datasource_names(datasources_dir: str) -> List[str]: +def _yaml_list_datasource_names(datasources_dir: str) -> list[str]: if not os.path.isdir(datasources_dir): return [] return [ @@ -162,7 +161,7 @@ def _sqlite_models_has_data_source_column(conn: sqlite3.Connection) -> bool: return any(r[1] == "data_source" for r in rows) -def _sqlite_list_datasource_names(conn: sqlite3.Connection) -> List[str]: +def _sqlite_list_datasource_names(conn: sqlite3.Connection) -> list[str]: cur = conn.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name='datasources'" ) @@ -191,10 +190,10 @@ def migrate_sqlite_schema(db_path: str) -> None: rows = conn.execute("SELECT name, data FROM models").fetchall() available = _sqlite_list_datasource_names(conn) - migrated: List[tuple] = [] + migrated: list[tuple] = [] for name, blob in rows: data = json.loads(blob) - ds: Optional[str] = data.get("data_source") or None + ds: str | None = data.get("data_source") or None if not ds: ds = _resolve_orphan_data_source(name=name, available_datasources=available) data["data_source"] = ds diff --git a/slayer/storage/v5_migration.py b/slayer/storage/v5_migration.py index 64762c78..7940010e 100644 --- a/slayer/storage/v5_migration.py +++ b/slayer/storage/v5_migration.py @@ -12,7 +12,6 @@ from __future__ import annotations -from typing import Optional from slayer.storage.migrations import register_migration @@ -32,7 +31,7 @@ }) -def _coerce_type_string(t: Optional[str]) -> tuple[bool, Optional[str]]: +def _coerce_type_string(t: str | None) -> tuple[bool, str | None]: """Map a single ``type`` value. Returns ``(should_set, new_value)``. * Legacy lowercase canonical → mapped uppercase. diff --git a/slayer/storage/yaml_storage.py b/slayer/storage/yaml_storage.py index d38740d7..91bfc413 100644 --- a/slayer/storage/yaml_storage.py +++ b/slayer/storage/yaml_storage.py @@ -23,7 +23,8 @@ import contextlib import os -from typing import Any, Dict, Iterator, List, Optional, Tuple +from typing import Any +from collections.abc import Iterator import yaml from pydantic import ValidationError @@ -112,8 +113,8 @@ async def _save_model_impl(self, model: SlayerModel) -> None: with open(path, "w") as f: yaml.dump(data, f, sort_keys=False) - async def _list_all_model_identities(self) -> List[Tuple[str, str]]: - result: List[Tuple[str, str]] = [] + async def _list_all_model_identities(self) -> list[tuple[str, str]]: + result: list[tuple[str, str]] = [] if not os.path.isdir(self.models_dir): return result for ds in sorted(os.listdir(self.models_dir)): @@ -128,8 +129,8 @@ async def _list_all_model_identities(self) -> List[Tuple[str, str]]: async def get_model( self, name: str, - data_source: Optional[str] = None, - ) -> Optional[SlayerModel]: + data_source: str | None = None, + ) -> SlayerModel | None: target = await self._resolve_target_or_none(name, data_source=data_source) if target is None: return None @@ -166,9 +167,9 @@ async def update_column_sampled( data_source: str, model_name: str, column_name: str, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: path = self._model_path(data_source, model_name) if not os.path.exists(path): @@ -204,7 +205,7 @@ async def save_datasource(self, datasource: DatasourceConfig) -> None: with open(path, "w") as f: yaml.dump(data, f, sort_keys=False) - async def get_datasource(self, name: str) -> Optional[DatasourceConfig]: + async def get_datasource(self, name: str) -> DatasourceConfig | None: # DEV-1405: sanitize before composing the filesystem path. _validate_path_component(name, kind="datasource name") path = os.path.join(self.datasources_dir, f"{name}.yaml") @@ -224,7 +225,7 @@ async def get_datasource(self, name: str) -> Optional[DatasourceConfig]: f"Datasource '{name}': invalid config — {exc}" ) from exc - async def list_datasources(self) -> List[str]: + async def list_datasources(self) -> list[str]: result = [] for filename in sorted(os.listdir(self.datasources_dir)): if filename.endswith((".yaml", ".yml")): @@ -240,7 +241,7 @@ async def _delete_datasource_row(self, name: str) -> bool: # ---- datasource priority ----------------------------------------------- - async def get_datasource_priority(self) -> List[str]: + async def get_datasource_priority(self) -> list[str]: if not os.path.exists(self._priority_path): return [] with open(self._priority_path) as f: # NOSONAR(S7493) — YAMLStorage uses sync I/O inside async by design (CLAUDE.md, Async Architecture) @@ -250,13 +251,13 @@ async def get_datasource_priority(self) -> List[str]: return [] return [str(p) for p in priority] - async def _set_datasource_priority_raw(self, priority: List[str]) -> None: + async def _set_datasource_priority_raw(self, priority: list[str]) -> None: with open(self._priority_path, "w") as f: # NOSONAR(S7493) — YAMLStorage uses sync I/O inside async by design (CLAUDE.md, Async Architecture) yaml.dump({"priority": list(priority)}, f, sort_keys=False) # ---- memories (DEV-1357 v2) ------------------------------------------- - def _read_yaml_list(self, path: str) -> List[Dict[str, Any]]: + def _read_yaml_list(self, path: str) -> list[dict[str, Any]]: if not os.path.exists(path): return [] with open(path) as f: # NOSONAR(S7493) — YAMLStorage uses sync I/O inside async by design (CLAUDE.md, Async Architecture) @@ -265,7 +266,7 @@ def _read_yaml_list(self, path: str) -> List[Dict[str, Any]]: return [] return [d for d in data if isinstance(d, dict)] - def _write_yaml_list(self, path: str, rows: List[Dict[str, Any]]) -> None: + def _write_yaml_list(self, path: str, rows: list[dict[str, Any]]) -> None: with open(path, "w") as f: # NOSONAR(S7493) — YAMLStorage uses sync I/O inside async by design (CLAUDE.md, Async Architecture) yaml.dump(rows, f, sort_keys=False) @@ -300,12 +301,12 @@ async def _next_memory_seq(self) -> str: return str(max_id + 1) def _normalize_legacy_rows( - self, rows: List[Dict[str, Any]], - ) -> List[Dict[str, Any]]: + self, rows: list[dict[str, Any]], + ) -> list[dict[str, Any]]: """DEV-1428: dedupe legacy duplicate rows where the same logical id exists in both int and str form. Fails loud when their content differs.""" - seen: Dict[str, Dict[str, Any]] = {} + seen: dict[str, dict[str, Any]] = {} for row in rows: raw = row.get("id") if isinstance(raw, bool): @@ -337,7 +338,7 @@ def _normalize_legacy_rows( return list(seen.values()) @staticmethod - def _rows_content_equal(a: Dict[str, Any], b: Dict[str, Any]) -> bool: + def _rows_content_equal(a: dict[str, Any], b: dict[str, Any]) -> bool: # DEV-1428: "content" excludes ``created_at`` — two legacy rows for # the same logical memory may carry different timestamps (e.g. one # written on int-id v1, then re-saved as str on v2). The plan's @@ -355,7 +356,7 @@ async def _save_memory_row(self, memory: Memory) -> None: rows.append(memory.model_dump(mode="json")) self._write_yaml_list(self._memories_path, rows) - async def _get_memory_row(self, memory_id: str) -> Optional[Memory]: + async def _get_memory_row(self, memory_id: str) -> Memory | None: rows = self._normalize_legacy_rows( self._read_yaml_list(self._memories_path), ) @@ -365,8 +366,8 @@ async def _get_memory_row(self, memory_id: str) -> Optional[Memory]: return None async def _list_memories_rows( - self, *, entities: Optional[List[str]] - ) -> List[Memory]: + self, *, entities: list[str] | None + ) -> list[Memory]: rows = self._normalize_legacy_rows( self._read_yaml_list(self._memories_path), ) diff --git a/tests/conftest.py b/tests/conftest.py index d89ba5c4..5e6f1220 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,7 +2,7 @@ import os import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest diff --git a/tests/facade/test_catalog.py b/tests/facade/test_catalog.py index 9aca303f..c7d3abc1 100644 --- a/tests/facade/test_catalog.py +++ b/tests/facade/test_catalog.py @@ -3,7 +3,6 @@ from __future__ import annotations import logging -from typing import List from slayer.core.enums import DataType from slayer.core.models import ( @@ -27,10 +26,10 @@ def _model( *, name: str, data_source: str = "ds1", - columns: List[Column] | None = None, - measures: List[ModelMeasure] | None = None, - aggregations: List[Aggregation] | None = None, - joins: List[ModelJoin] | None = None, + columns: list[Column] | None = None, + measures: list[ModelMeasure] | None = None, + aggregations: list[Aggregation] | None = None, + joins: list[ModelJoin] | None = None, hidden: bool = False, sql: str | None = None, description: str | None = None, diff --git a/tests/facade/test_catalog_flatten_filter.py b/tests/facade/test_catalog_flatten_filter.py index 751771c0..ac69a236 100644 --- a/tests/facade/test_catalog_flatten_filter.py +++ b/tests/facade/test_catalog_flatten_filter.py @@ -24,7 +24,6 @@ from __future__ import annotations -from typing import List from slayer.core.enums import DataType from slayer.core.models import ( @@ -46,10 +45,10 @@ def _model( *, name: str, - columns: List[Column], - joins: List[ModelJoin] | None = None, - measures: List[ModelMeasure] | None = None, - aggregations: List[Aggregation] | None = None, + columns: list[Column], + joins: list[ModelJoin] | None = None, + measures: list[ModelMeasure] | None = None, + aggregations: list[Aggregation] | None = None, ) -> SlayerModel: return SlayerModel( name=name, diff --git a/tests/facade/test_catalog_sql_metabase_golden.py b/tests/facade/test_catalog_sql_metabase_golden.py index a95766a8..4a03f903 100644 --- a/tests/facade/test_catalog_sql_metabase_golden.py +++ b/tests/facade/test_catalog_sql_metabase_golden.py @@ -16,7 +16,6 @@ from __future__ import annotations from pathlib import Path -from typing import List import pytest @@ -104,11 +103,11 @@ def _translate(sql: str, catalog: FacadeCatalog): # --- corpus loading --------------------------------------------------------- -def _load_corpus_statements() -> List[str]: +def _load_corpus_statements() -> list[str]: """Yield each ``-- #N ...`` block from the captured corpus as a single string.""" text = FIXTURE.read_text(encoding="utf-8") - statements: List[str] = [] - current: List[str] = [] + statements: list[str] = [] + current: list[str] = [] capturing = False for line in text.splitlines(): if line.startswith("-- #") and " first seen at " in line: diff --git a/tests/flight/conftest.py b/tests/flight/conftest.py index 01f4f7ba..337165c1 100644 --- a/tests/flight/conftest.py +++ b/tests/flight/conftest.py @@ -18,7 +18,7 @@ import time import urllib.request from pathlib import Path -from typing import Callable, Iterator, Tuple +from collections.abc import Callable, Iterator import pytest @@ -121,7 +121,7 @@ def _connect(url: str, driver_args: list[str] | None = None): @pytest.fixture -def capture_stub(tmp_path: Path) -> Iterator[Tuple[str, Path]]: +def capture_stub(tmp_path: Path) -> Iterator[tuple[str, Path]]: """Spin up a CaptureFlightServer on an ephemeral port. Yields ``(grpc_location, log_path)`` where ``log_path`` accumulates one diff --git a/tests/integration/_pg_serve_helpers.py b/tests/integration/_pg_serve_helpers.py index e2ad99f5..c3a0c035 100644 --- a/tests/integration/_pg_serve_helpers.py +++ b/tests/integration/_pg_serve_helpers.py @@ -15,7 +15,6 @@ import tempfile import threading import time -from typing import List, Optional, Tuple import pytest @@ -29,11 +28,11 @@ def start_pg_demo_server( *, - token: Optional[str], - log_records: Optional[List[logging.LogRecord]] = None, - storage_sink: Optional[list] = None, + token: str | None, + log_records: list[logging.LogRecord] | None = None, + storage_sink: list | None = None, bind_host: str = "127.0.0.1", -) -> Tuple[asyncio.AbstractEventLoop, threading.Thread, str, int]: +) -> tuple[asyncio.AbstractEventLoop, threading.Thread, str, int]: """Boot a Postgres-facade server backed by the Jaffle Shop demo. Returns ``(loop, thread, host, port)``. Caller stops the server via @@ -132,7 +131,7 @@ async def _setup(): class _ListHandler(logging.Handler): - def __init__(self, sink: List[logging.LogRecord]) -> None: + def __init__(self, sink: list[logging.LogRecord]) -> None: super().__init__() self._sink = sink diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 4ba0df42..a47d9fd4 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -19,7 +19,8 @@ import time import urllib.request from pathlib import Path -from typing import Any, Callable, Iterator, Optional, Tuple +from typing import Any +from collections.abc import Callable, Iterator import pytest @@ -80,8 +81,8 @@ def _format_flight_jdbc_url( host: str, port: int, use_encryption: bool = False, - token: Optional[str] = None, - environment_id: Optional[str] = None, + token: str | None = None, + environment_id: str | None = None, ) -> str: params = [f"useEncryption={'true' if use_encryption else 'false'}"] if token is not None: @@ -172,7 +173,7 @@ def _connect(url: str, driver_args: list[str] | None = None): return _connect -def _start_flight_demo_server(*, token: Optional[str]): +def _start_flight_demo_server(*, token: str | None): """Boot a Flight SQL server backed by the bundled Jaffle Shop demo. Returns ``(server, host, port)``. The caller is responsible for @@ -204,7 +205,7 @@ def _start_flight_demo_server(*, token: Optional[str]): @pytest.fixture(scope="module") -def flight_demo_server() -> Iterator[Tuple[str, int]]: +def flight_demo_server() -> Iterator[tuple[str, int]]: """Yield ``(host, port)`` of a no-auth Flight SQL server backed by the Jaffle Shop demo.""" server, host, port = _start_flight_demo_server(token=None) try: @@ -215,7 +216,7 @@ def flight_demo_server() -> Iterator[Tuple[str, int]]: @pytest.fixture(scope="module") -def flight_demo_server_with_token() -> Iterator[Tuple[str, int, str]]: +def flight_demo_server_with_token() -> Iterator[tuple[str, int, str]]: """Same as ``flight_demo_server`` but with a bearer token enforced.""" token = "s3cret" server, host, port = _start_flight_demo_server(token=token) diff --git a/tests/integration/conftest_metabase.py b/tests/integration/conftest_metabase.py index 4fe001ec..be261fcb 100644 --- a/tests/integration/conftest_metabase.py +++ b/tests/integration/conftest_metabase.py @@ -21,7 +21,8 @@ import subprocess import time import uuid -from typing import Any, Dict, Iterator, List, Optional, Tuple +from typing import Any +from collections.abc import Iterator import pytest import requests @@ -86,22 +87,22 @@ def _raise_with_body(resp: "requests.Response", path: str) -> None: response=resp, ) - def post(self, path: str, json_body: Optional[dict] = None, *, timeout: int = 30) -> dict: + def post(self, path: str, json_body: dict | None = None, *, timeout: int = 30) -> dict: resp = self._session.post(self._url(path), json=json_body or {}, timeout=timeout) self._raise_with_body(resp, path) return resp.json() if resp.content else {} - def get(self, path: str, params: Optional[dict] = None, *, timeout: int = 30) -> Any: + def get(self, path: str, params: dict | None = None, *, timeout: int = 30) -> Any: resp = self._session.get(self._url(path), params=params, timeout=timeout) self._raise_with_body(resp, path) return resp.json() if resp.content else {} - def put(self, path: str, json_body: Optional[dict] = None, *, timeout: int = 30) -> dict: + def put(self, path: str, json_body: dict | None = None, *, timeout: int = 30) -> dict: resp = self._session.put(self._url(path), json=json_body or {}, timeout=timeout) self._raise_with_body(resp, path) return resp.json() if resp.content else {} - def post_raw(self, path: str, json_body: Optional[dict] = None, *, timeout: int = 30): + def post_raw(self, path: str, json_body: dict | None = None, *, timeout: int = 30): """POST without raise_for_status — for error-envelope tests.""" return self._session.post(self._url(path), json=json_body or {}, timeout=timeout) @@ -160,10 +161,10 @@ class MetabaseE2EEnv(BaseModel): session_token: str client: MetabaseClient token_db_id: int - pg_primary: Tuple[str, int] + pg_primary: tuple[str, int] pg_primary_password: str - pg_auth: Tuple[str, int, str] - log_records: List[logging.LogRecord] + pg_auth: tuple[str, int, str] + log_records: list[logging.LogRecord] pg_primary_storage: Any def make_client(self, db_id: int) -> MetabaseClient: @@ -193,7 +194,7 @@ def _docker_available() -> bool: return result.returncode == 0 and result.stdout.strip() != "" -def _run_metabase_container() -> Tuple[str, int]: +def _run_metabase_container() -> tuple[str, int]: """Start the Metabase container; return ``(container_id, host_port)``.""" container_name = f"slayer-mb-e2e-{uuid.uuid4().hex[:8]}" cmd = [ @@ -238,7 +239,7 @@ def _run_metabase_container() -> Tuple[str, int]: return container_id, host_port -def _stop_container(container_id: Optional[str]) -> None: +def _stop_container(container_id: str | None) -> None: if not container_id: return try: @@ -282,7 +283,7 @@ def _wait_for_health(base_url: str, timeout_s: int) -> bool: return False -def _fetch_setup_token(base_url: str) -> Optional[str]: +def _fetch_setup_token(base_url: str) -> str | None: r = requests.get(f"{base_url}/api/session/properties", timeout=10) r.raise_for_status() props = r.json() @@ -376,10 +377,10 @@ def _register_database( def _wait_for_metadata( base_url: str, session_token: str, db_id: int, *, min_tables: int, timeout_s: int -) -> Dict[str, Any]: +) -> dict[str, Any]: deadline = time.monotonic() + timeout_s headers = {"X-Metabase-Session": session_token} - last_payload: Dict[str, Any] = {} + last_payload: dict[str, Any] = {} while time.monotonic() < deadline: r = requests.get( f"{base_url}/api/database/{db_id}/metadata", @@ -405,7 +406,7 @@ def _wait_for_metadata( def _bootstrap_metabase_session( *, container_id: str, host_port: int, port_a: int, port_b: int, -) -> Tuple[str, str, int, int]: +) -> tuple[str, str, int, int]: """Boot Metabase, run first-time setup, register both pg-serve databases. Returns ``(base_url, session_token, db_id_no_token, db_id_token)``. @@ -460,11 +461,11 @@ def metabase_e2e_env() -> Iterator[MetabaseE2EEnv]: if not _docker_available(): pytest.skip("Docker is unavailable; the metabase_e2e suite needs a working Docker daemon") - log_records: List[logging.LogRecord] = [] + log_records: list[logging.LogRecord] = [] storage_sink: list = [] loop_a = thread_a = None loop_b = thread_b = None - container_id: Optional[str] = None + container_id: str | None = None try: loop_a, thread_a, host_a, port_a = start_pg_demo_server( @@ -521,14 +522,14 @@ def metabase_e2e_env() -> Iterator[MetabaseE2EEnv]: # Used by tests that re-emit the env in a per-test mutation context. -def encode_native_query(sql: str) -> Dict[str, Any]: +def encode_native_query(sql: str) -> dict[str, Any]: """Convenience helper used by native-SQL tests (E.7, G.4, etc.).""" return {"type": "native", "native": {"query": sql, "template-tags": {}}} -def encode_mbql_query(*, source_table: int, **extras: Any) -> Dict[str, Any]: +def encode_mbql_query(*, source_table: int, **extras: Any) -> dict[str, Any]: """Build an MBQL ``query`` body for ``/api/dataset``.""" - inner: Dict[str, Any] = {"source-table": source_table} + inner: dict[str, Any] = {"source-table": source_table} inner.update(extras) return {"type": "query", "query": inner} diff --git a/tests/integration/test_in_memory_sqlite.py b/tests/integration/test_in_memory_sqlite.py index 8d67eb11..551bcdb4 100644 --- a/tests/integration/test_in_memory_sqlite.py +++ b/tests/integration/test_in_memory_sqlite.py @@ -9,7 +9,7 @@ """ from pathlib import Path -from typing import Iterable +from collections.abc import Iterable import pytest import sqlalchemy as sa diff --git a/tests/integration/test_ingestion_jaffle_shop.py b/tests/integration/test_ingestion_jaffle_shop.py index 5505757b..d4d30683 100644 --- a/tests/integration/test_ingestion_jaffle_shop.py +++ b/tests/integration/test_ingestion_jaffle_shop.py @@ -13,7 +13,6 @@ import shutil import tempfile from pathlib import Path -from typing import Dict import pytest import yaml @@ -86,7 +85,7 @@ def jaffle_models(jaffle_duckdb_path): # ``column_name -> (DataType, primary_key, format_type_or_None)``. Columns # whose source name contained a dot (joined columns) are excluded by # ``_columns_to_model`` and so are not listed here. -_EXPECTED_COLUMNS: Dict[str, Dict[str, tuple]] = { +_EXPECTED_COLUMNS: dict[str, dict[str, tuple]] = { "customers": { "id": (DataType.TEXT, True, None), "name": (DataType.TEXT, False, None), @@ -135,7 +134,7 @@ def jaffle_models(jaffle_duckdb_path): # Expected direct joins per model. Values are sets of (target_model, src_col, tgt_col). -_EXPECTED_JOINS: Dict[str, set] = { +_EXPECTED_JOINS: dict[str, set] = { "customers": set(), "stores": set(), "products": set(), diff --git a/tests/integration/test_integration_flight.py b/tests/integration/test_integration_flight.py index 6f907428..629953ef 100644 --- a/tests/integration/test_integration_flight.py +++ b/tests/integration/test_integration_flight.py @@ -14,7 +14,8 @@ from __future__ import annotations import threading -from typing import Any, Callable, Tuple +from typing import Any +from collections.abc import Callable import pytest @@ -62,7 +63,7 @@ def _columns_of(rs) -> list[str]: def test_get_catalogs( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -84,7 +85,7 @@ def test_get_catalogs( def test_get_schemas_and_tables( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -113,7 +114,7 @@ def test_get_schemas_and_tables( def test_get_primary_keys_empty( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -134,7 +135,7 @@ def test_get_primary_keys_empty( def test_get_keys_and_cross_reference_empty( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -152,7 +153,7 @@ def test_get_keys_and_cross_reference_empty( def test_get_type_info_returns_jdbc_shape( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -180,7 +181,7 @@ def test_get_type_info_returns_jdbc_shape( def test_select_information_schema_metrics( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -202,7 +203,7 @@ def test_select_information_schema_metrics( def test_select_row_count_via_prepared_statement( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -220,7 +221,7 @@ def test_select_row_count_via_prepared_statement( def test_time_grain_select( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -243,7 +244,7 @@ def test_time_grain_select( def test_cross_model_dim_select( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -268,7 +269,7 @@ def test_cross_model_dim_select( def test_select_star_rejected( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -283,7 +284,7 @@ def test_select_star_rejected( def test_dml_rejected( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -310,7 +311,7 @@ def test_dml_rejected( ], ) def test_probe_queries( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], probe_sql: str, @@ -352,7 +353,7 @@ def test_probe_queries( ), ) def test_auth_positive( - flight_demo_server_with_token: Tuple[str, int, str], + flight_demo_server_with_token: tuple[str, int, str], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -372,7 +373,7 @@ def test_auth_positive( def test_auth_negative_wrong_token( - flight_demo_server_with_token: Tuple[str, int, str], + flight_demo_server_with_token: tuple[str, int, str], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: @@ -402,7 +403,7 @@ def test_auth_negative_wrong_token( def test_n10_concurrent_executequery( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], flight_jdbc_url: Callable[..., str], jaydebeapi_connect: Callable[..., Any], ) -> None: diff --git a/tests/integration/test_integration_flight_pyarrow_client.py b/tests/integration/test_integration_flight_pyarrow_client.py index 042a6a14..37ed27be 100644 --- a/tests/integration/test_integration_flight_pyarrow_client.py +++ b/tests/integration/test_integration_flight_pyarrow_client.py @@ -15,7 +15,6 @@ from __future__ import annotations import threading -from typing import Tuple import pyarrow.flight as fl import pytest @@ -61,7 +60,7 @@ def _bearer_options(token: str | None) -> fl.FlightCallOptions: # ----- catalog commands ------------------------------------------------------ -def test_get_catalogs(flight_demo_server: Tuple[str, int]) -> None: +def test_get_catalogs(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) descriptor = _descriptor_for( @@ -73,7 +72,7 @@ def test_get_catalogs(flight_demo_server: Tuple[str, int]) -> None: assert any(r["catalog_name"] == "slayer" for r in rows) -def test_get_db_schemas(flight_demo_server: Tuple[str, int]) -> None: +def test_get_db_schemas(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) descriptor = _descriptor_for( @@ -85,7 +84,7 @@ def test_get_db_schemas(flight_demo_server: Tuple[str, int]) -> None: assert "jaffle_shop" in names -def test_get_tables(flight_demo_server: Tuple[str, int]) -> None: +def test_get_tables(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) descriptor = _descriptor_for(fsql_pb.CommandGetTables(), "CommandGetTables") @@ -97,7 +96,7 @@ def test_get_tables(flight_demo_server: Tuple[str, int]) -> None: assert ("jaffle_shop", "customers") in pairs -def test_get_table_types(flight_demo_server: Tuple[str, int]) -> None: +def test_get_table_types(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) descriptor = _descriptor_for( @@ -109,7 +108,7 @@ def test_get_table_types(flight_demo_server: Tuple[str, int]) -> None: assert {"TABLE", "VIEW"} <= rows -def test_get_primary_keys_empty(flight_demo_server: Tuple[str, int]) -> None: +def test_get_primary_keys_empty(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) cmd = fsql_pb.CommandGetPrimaryKeys() @@ -123,7 +122,7 @@ def test_get_primary_keys_empty(flight_demo_server: Tuple[str, int]) -> None: assert "column_name" in table.schema.names -def test_get_sql_info(flight_demo_server: Tuple[str, int]) -> None: +def test_get_sql_info(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) descriptor = _descriptor_for(fsql_pb.CommandGetSqlInfo(), "CommandGetSqlInfo") @@ -163,7 +162,7 @@ def _execute_prepared(client: fl.FlightClient, handle: bytes): return client.do_get(info.endpoints[0].ticket).read_all() -def test_prepared_statement_row_count(flight_demo_server: Tuple[str, int]) -> None: +def test_prepared_statement_row_count(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) resp = _create_prepared(client, "SELECT row_count FROM orders") @@ -174,7 +173,7 @@ def test_prepared_statement_row_count(flight_demo_server: Tuple[str, int]) -> No assert int(table.to_pylist()[0]["row_count"]) > 0 -def test_prepared_statement_time_grain(flight_demo_server: Tuple[str, int]) -> None: +def test_prepared_statement_time_grain(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) sql = ( @@ -190,7 +189,7 @@ def test_prepared_statement_time_grain(flight_demo_server: Tuple[str, int]) -> N assert int(row["row_count"]) > 0 -def test_prepared_statement_cross_model_dim(flight_demo_server: Tuple[str, int]) -> None: +def test_prepared_statement_cross_model_dim(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) resp = _create_prepared( @@ -202,7 +201,7 @@ def test_prepared_statement_cross_model_dim(flight_demo_server: Tuple[str, int]) def test_prepared_statement_info_schema_metrics( - flight_demo_server: Tuple[str, int], + flight_demo_server: tuple[str, int], ) -> None: host, port = flight_demo_server client = _client(host=host, port=port) @@ -222,7 +221,7 @@ def test_prepared_statement_info_schema_metrics( ], ) def test_prepared_statement_probe_queries( - flight_demo_server: Tuple[str, int], probe_sql: str, + flight_demo_server: tuple[str, int], probe_sql: str, ) -> None: host, port = flight_demo_server client = _client(host=host, port=port) @@ -238,7 +237,7 @@ def test_prepared_statement_probe_queries( # ----- error paths ----------------------------------------------------------- -def test_select_star_rejected(flight_demo_server: Tuple[str, int]) -> None: +def test_select_star_rejected(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) with pytest.raises(fl.FlightServerError) as excinfo: @@ -246,7 +245,7 @@ def test_select_star_rejected(flight_demo_server: Tuple[str, int]) -> None: assert "SELECT * not supported" in str(excinfo.value) -def test_dml_rejected(flight_demo_server: Tuple[str, int]) -> None: +def test_dml_rejected(flight_demo_server: tuple[str, int]) -> None: host, port = flight_demo_server client = _client(host=host, port=port) with pytest.raises(fl.FlightServerError) as excinfo: @@ -254,7 +253,7 @@ def test_dml_rejected(flight_demo_server: Tuple[str, int]) -> None: assert "read-only" in str(excinfo.value).lower() -def test_close_prepared_statement(flight_demo_server: Tuple[str, int]) -> None: +def test_close_prepared_statement(flight_demo_server: tuple[str, int]) -> None: """``ActionClosePreparedStatementRequest`` is a no-op; it must complete cleanly.""" host, port = flight_demo_server client = _client(host=host, port=port) @@ -270,7 +269,7 @@ def test_close_prepared_statement(flight_demo_server: Tuple[str, int]) -> None: def test_auth_positive( - flight_demo_server_with_token: Tuple[str, int, str], + flight_demo_server_with_token: tuple[str, int, str], ) -> None: """With the correct bearer token attached on every RPC, the server accepts.""" host, port, token = flight_demo_server_with_token @@ -285,7 +284,7 @@ def test_auth_positive( def test_auth_negative_missing_token( - flight_demo_server_with_token: Tuple[str, int, str], + flight_demo_server_with_token: tuple[str, int, str], ) -> None: """Without an Authorization header the server rejects with UNAUTHENTICATED.""" host, port, _token = flight_demo_server_with_token @@ -296,7 +295,7 @@ def test_auth_negative_missing_token( def test_auth_negative_wrong_token( - flight_demo_server_with_token: Tuple[str, int, str], + flight_demo_server_with_token: tuple[str, int, str], ) -> None: host, port, _token = flight_demo_server_with_token client = _client(host=host, port=port) @@ -308,7 +307,7 @@ def test_auth_negative_wrong_token( # ----- concurrency ----------------------------------------------------------- -def test_n10_concurrent_prepared_statements(flight_demo_server: Tuple[str, int]) -> None: +def test_n10_concurrent_prepared_statements(flight_demo_server: tuple[str, int]) -> None: """Ten parallel prepared-statement round-trips against the same server.""" host, port = flight_demo_server diff --git a/tests/integration/test_integration_pg_facade.py b/tests/integration/test_integration_pg_facade.py index 3ce314e4..3a76cd83 100644 --- a/tests/integration/test_integration_pg_facade.py +++ b/tests/integration/test_integration_pg_facade.py @@ -9,7 +9,7 @@ from __future__ import annotations import asyncio -from typing import Iterator, Tuple +from collections.abc import Iterator import pytest @@ -21,7 +21,7 @@ @pytest.fixture(scope="module") -def pg_demo_server() -> Iterator[Tuple[str, int]]: +def pg_demo_server() -> Iterator[tuple[str, int]]: loop, thread, host, port = start_pg_demo_server(token=None) try: yield host, port @@ -31,7 +31,7 @@ def pg_demo_server() -> Iterator[Tuple[str, int]]: @pytest.fixture(scope="module") -def pg_demo_server_with_token() -> Iterator[Tuple[str, int, str]]: +def pg_demo_server_with_token() -> Iterator[tuple[str, int, str]]: token = "s3cret" loop, thread, host, port = start_pg_demo_server(token=token) try: diff --git a/tests/integration/test_mcp_inspect.py b/tests/integration/test_mcp_inspect.py index 0cab4e8f..9257f1ac 100644 --- a/tests/integration/test_mcp_inspect.py +++ b/tests/integration/test_mcp_inspect.py @@ -8,7 +8,7 @@ """ import sqlite3 -from typing import Any, Optional +from typing import Any import pytest @@ -257,7 +257,7 @@ class TestInspectModelEndToEnd: """Full ``inspect_model`` run against the SQLite fixture — confirms every section of the markdown output is produced end-to-end.""" - async def _call(self, server: Any, *, name: str, arguments: Optional[dict] = None) -> str: + async def _call(self, server: Any, *, name: str, arguments: dict | None = None) -> str: content, _ = await server.call_tool(name=name, arguments=arguments or {}) return content[0].text @@ -333,7 +333,7 @@ class TestInspectModelSectionGatingIntegration: short-circuits actually take effect when sections are dropped. """ - async def _call(self, server: Any, *, name: str, arguments: Optional[dict] = None) -> str: + async def _call(self, server: Any, *, name: str, arguments: dict | None = None) -> str: content, _ = await server.call_tool(name=name, arguments=arguments or {}) return content[0].text diff --git a/tests/integration/test_metabase_e2e.py b/tests/integration/test_metabase_e2e.py index 27e550a0..d55d86cd 100644 --- a/tests/integration/test_metabase_e2e.py +++ b/tests/integration/test_metabase_e2e.py @@ -17,7 +17,7 @@ import logging import math import time -from typing import Any, Dict, List, Tuple +from typing import Any import pytest @@ -62,11 +62,11 @@ async def _scalar(host: str, port: int, password: str, sql: str) -> Any: await conn.close() -def _dataset_rows(payload: Dict[str, Any]) -> List[List[Any]]: +def _dataset_rows(payload: dict[str, Any]) -> list[list[Any]]: return payload.get("data", {}).get("rows", []) or [] -def _dataset_cols(payload: Dict[str, Any]) -> List[Dict[str, Any]]: +def _dataset_cols(payload: dict[str, Any]) -> list[dict[str, Any]]: return payload.get("data", {}).get("cols", []) or [] @@ -213,7 +213,7 @@ def _wait_until(predicate, *, timeout_s: float = 20, interval_s: float = 0.5) -> def _wait_until_stable( getter, *, timeout_s: float = 20, settle_s: float = 2.0, interval_s: float = 0.5, -) -> Tuple[bool, Any]: +) -> tuple[bool, Any]: """Poll ``getter()`` until its value stops changing for ``settle_s`` seconds. Returns ``(stabilized, last_value)`` so callers can distinguish a real settle from a timeout — silently returning the @@ -1273,7 +1273,7 @@ def run_one(tid: int) -> int: async def test_asyncpg_concurrent_connections(metabase_e2e_env: MetabaseE2EEnv) -> None: host, port = metabase_e2e_env.pg_primary - async def one(idx: int) -> Tuple[int, str]: + async def one(idx: int) -> tuple[int, str]: conn = await _asyncpg_connect(host, port, password=metabase_e2e_env.pg_primary_password) try: # Each connection sets its own application_name marker so we can diff --git a/tests/pg_facade/test_connection.py b/tests/pg_facade/test_connection.py index 6f380810..c9ddeb37 100644 --- a/tests/pg_facade/test_connection.py +++ b/tests/pg_facade/test_connection.py @@ -8,7 +8,6 @@ import asyncio import struct import types -from typing import List, Optional, Tuple import pytest @@ -45,10 +44,10 @@ class _FakeStorage: def __init__(self, models_by_ds) -> None: self._models_by_ds = models_by_ds - async def list_datasources(self) -> List[str]: # NOSONAR(S7503) — async to satisfy the awaited interface + async def list_datasources(self) -> list[str]: # NOSONAR(S7503) — async to satisfy the awaited interface return list(self._models_by_ds) - async def list_models(self, *, data_source: str) -> List[str]: # NOSONAR(S7503) — async to satisfy the awaited interface + async def list_models(self, *, data_source: str) -> list[str]: # NOSONAR(S7503) — async to satisfy the awaited interface return [m.name for m in self._models_by_ds.get(data_source, [])] async def get_model(self, *, name: str, data_source: str): # NOSONAR(S7503) — async to satisfy the awaited interface @@ -141,7 +140,7 @@ def _terminate() -> bytes: return _frame(b"X", b"") -def _parse(name: str, sql: str, oids: Tuple[int, ...] = ()) -> bytes: +def _parse(name: str, sql: str, oids: tuple[int, ...] = ()) -> bytes: body = name.encode() + b"\x00" + sql.encode() + b"\x00" + struct.pack(">h", len(oids)) for o in oids: body += struct.pack(">i", o) @@ -150,9 +149,9 @@ def _parse(name: str, sql: str, oids: Tuple[int, ...] = ()) -> bytes: def _bind( portal: str, stmt: str, *, - values: Tuple[Optional[bytes], ...] = (), - param_formats: Tuple[int, ...] = (), - result_formats: Tuple[int, ...] = (), + values: tuple[bytes | None, ...] = (), + param_formats: tuple[int, ...] = (), + result_formats: tuple[int, ...] = (), ) -> bytes: body = portal.encode() + b"\x00" + stmt.encode() + b"\x00" body += struct.pack(">h", len(param_formats)) @@ -190,7 +189,7 @@ def _close(kind: str, name: str) -> bytes: async def _run( - input_bytes: bytes, *, token: Optional[str] = None, storage=None, engine=None, + input_bytes: bytes, *, token: str | None = None, storage=None, engine=None, tls_ctx=None, ) -> _FakeWriter: reader = asyncio.StreamReader() @@ -208,19 +207,19 @@ async def _run( return writer -def _messages(buf: bytes, *, leading_raw: int = 0) -> List[Tuple[str, bytes]]: +def _messages(buf: bytes, *, leading_raw: int = 0) -> list[tuple[str, bytes]]: return proto.split_messages(bytes(buf[leading_raw:])) -def _types(msgs: List[Tuple[str, bytes]]) -> List[str]: +def _types(msgs: list[tuple[str, bytes]]) -> list[str]: return [t for t, _ in msgs] -def _ready_statuses(msgs: List[Tuple[str, bytes]]) -> List[bytes]: +def _ready_statuses(msgs: list[tuple[str, bytes]]) -> list[bytes]: return [body for t, body in msgs if t == "Z"] -def _error_sqlstate(body: bytes) -> Optional[str]: +def _error_sqlstate(body: bytes) -> str | None: i = 0 while i < len(body) and body[i:i + 1] != b"\x00": ftype = body[i:i + 1] @@ -1003,7 +1002,7 @@ async def test_simple_query_catalog_union_routes_to_executor() -> None: # --- DEV-1569: per-connection SET state --- -def _show_value(msgs: List[Tuple[str, bytes]]) -> str: +def _show_value(msgs: list[tuple[str, bytes]]) -> str: """Pull the single-text-column value out of the last DataRow in `msgs`. The format is: ``count(i16) | length(i32) | bytes``. SHOW always returns @@ -1014,11 +1013,11 @@ def _show_value(msgs: List[Tuple[str, bytes]]) -> str: return body[6:6 + length].decode() -def _parameter_status(msgs: List[Tuple[str, bytes]]) -> List[Tuple[str, str]]: +def _parameter_status(msgs: list[tuple[str, bytes]]) -> list[tuple[str, str]]: """Decode every ParameterStatus message (``S``) in the stream into ``[(name, value), …]``. Note that startup ParameterStatus burst messages are included too — callers filter by name.""" - out: List[Tuple[str, str]] = [] + out: list[tuple[str, str]] = [] for t, body in msgs: if t != "S": continue @@ -1029,8 +1028,8 @@ def _parameter_status(msgs: List[Tuple[str, bytes]]) -> List[Tuple[str, str]]: return out -def _command_complete_tags(msgs: List[Tuple[str, bytes]]) -> List[str]: - out: List[str] = [] +def _command_complete_tags(msgs: list[tuple[str, bytes]]) -> list[str]: + out: list[str] = [] for t, body in msgs: if t != "C": continue @@ -1501,7 +1500,7 @@ async def test_two_connections_isolated_under_set_config() -> None: # --- DEV-1569 Codex round 2 follow-ups --- -def _index_of_first(msgs: List[Tuple[str, bytes]], pred) -> int: +def _index_of_first(msgs: list[tuple[str, bytes]], pred) -> int: for i, (t, body) in enumerate(msgs): if pred(t, body): return i @@ -1781,8 +1780,8 @@ async def test_set_pushes_parameter_status_for_all_reportable_names( async def _run_capturing( - input_bytes: bytes, *, token: Optional[str] = None, storage=None, engine=None, -) -> Tuple[_FakeWriter, PgConnection]: + input_bytes: bytes, *, token: str | None = None, storage=None, engine=None, +) -> tuple[_FakeWriter, PgConnection]: """Variant of ``_run`` that returns both the writer and the connection so tests can introspect ``conn._portals`` post-Bind.""" reader = asyncio.StreamReader() diff --git a/tests/pg_facade/test_probes.py b/tests/pg_facade/test_probes.py index d228e166..d4d62b66 100644 --- a/tests/pg_facade/test_probes.py +++ b/tests/pg_facade/test_probes.py @@ -2,7 +2,6 @@ from __future__ import annotations -from typing import Dict import sqlglot @@ -220,7 +219,7 @@ def test_set_config_returns_value_and_signals_no_mutation_when_no_dict() -> None """When session_settings is None, set_config still returns the requested value via the existing match_pg_probe path AND does NOT mutate the module-level seed. Pin the no-shared-state invariant.""" - seed_snapshot: Dict[str, str] = dict(SESSION_SETTING_SEED) + seed_snapshot: dict[str, str] = dict(SESSION_SETTING_SEED) batch = _probe("SELECT set_config('application_name', 'foo', false)") assert batch is not None assert batch.rows == [{"set_config": "foo"}] @@ -232,7 +231,7 @@ def test_match_pg_probe_does_not_mutate_session_settings_for_set_config() -> Non for set_config — that's reserved for the connection's Execute path via `match_pg_probe_with_mutation`. This pins the Describe-phase purity contract.""" - s: Dict[str, str] = {"application_name": "before"} + s: dict[str, str] = {"application_name": "before"} batch = _probe( "SELECT set_config('application_name', 'after', false)", session_settings=s, @@ -278,7 +277,7 @@ def test_match_pg_probe_with_mutation_does_not_mutate_dict() -> None: mutation is conveyed via the returned SetSettingOp and applied later by the connection. This keeps the Describe-phase pure even if it incidentally calls match_pg_probe_with_mutation.""" - s: Dict[str, str] = {"application_name": "before"} + s: dict[str, str] = {"application_name": "before"} outcome = match_pg_probe_with_mutation( _parse("SELECT set_config('application_name', 'after', false)"), datasource="jaffle", version_str="x", @@ -415,7 +414,7 @@ def test_match_pg_probe_with_mutation_no_dict_returns_outcome_for_set_config() - set_config. The mutation hint is what callers want; whether they have a dict to apply it to is up to them (and SESSION_SETTING_SEED must stay untouched).""" - seed_snapshot: Dict[str, str] = dict(SESSION_SETTING_SEED) + seed_snapshot: dict[str, str] = dict(SESSION_SETTING_SEED) outcome = match_pg_probe_with_mutation( _parse("SELECT set_config('application_name', 'val', false)"), datasource="jaffle", version_str="x", session_settings=None, diff --git a/tests/test_async_engine_disposal.py b/tests/test_async_engine_disposal.py index df124d0f..2f96ab59 100644 --- a/tests/test_async_engine_disposal.py +++ b/tests/test_async_engine_disposal.py @@ -7,7 +7,8 @@ import math import tempfile from pathlib import Path -from typing import Any, Iterator +from typing import Any +from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock import pytest diff --git a/tests/test_bm25_retriever.py b/tests/test_bm25_retriever.py index 77213321..8b99ef87 100644 --- a/tests/test_bm25_retriever.py +++ b/tests/test_bm25_retriever.py @@ -14,7 +14,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest_asyncio diff --git a/tests/test_cli_memory_string_ids.py b/tests/test_cli_memory_string_ids.py index e9b37c06..99e536b6 100644 --- a/tests/test_cli_memory_string_ids.py +++ b/tests/test_cli_memory_string_ids.py @@ -14,7 +14,7 @@ import sys import tempfile from contextlib import contextmanager -from typing import Iterator +from collections.abc import Iterator import pytest diff --git a/tests/test_cli_search_text_output.py b/tests/test_cli_search_text_output.py index 1d74127f..cce68741 100644 --- a/tests/test_cli_search_text_output.py +++ b/tests/test_cli_search_text_output.py @@ -8,7 +8,6 @@ import contextlib import io -from typing import List, Optional import pytest @@ -30,7 +29,7 @@ def _hit( id_: str = "1", score: float = 0.5, text: str = "", - description: Optional[str] = None, + description: str | None = None, ) -> SearchHit: return SearchHit( kind=kind, id=id_, score=score, text=text, description=description, @@ -92,7 +91,7 @@ def test_text_output_handles_multiple_hits() -> None: def test_text_output_takes_first_line_only(preview: str) -> None: response = SearchResponse(results=[_hit(description=preview)]) output = _capture_search_print(response) - lines: List[str] = output.splitlines() + lines: list[str] = output.splitlines() preview_line = next( ln.strip() for ln in lines if ln.startswith(" ") ) diff --git a/tests/test_client.py b/tests/test_client.py index 825e010f..afff7d9d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -10,7 +10,8 @@ import tempfile from types import MappingProxyType -from typing import Any, Dict, List, Mapping, Optional, Tuple +from typing import Any +from collections.abc import Mapping import pytest @@ -63,7 +64,7 @@ async def _save_orders(storage: YAMLStorage) -> None: # Minimal response payload that ``SlayerClient._parse_response`` consumes # without complaining. Used by the HTTP-mode mocks. -_CANNED_RESP: Dict[str, Any] = { +_CANNED_RESP: dict[str, Any] = { "data": [], "columns": [], "sql": "SELECT 1", @@ -80,16 +81,16 @@ class _CapturedRequests: """ def __init__(self) -> None: - self.calls: List[Dict[str, Any]] = [] + self.calls: list[dict[str, Any]] = [] def _replace_sync( self, *, method: str, path: str, - json: Optional[Dict[str, Any]] = None, - params: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: + json: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + ) -> dict[str, Any]: self.calls.append( {"method": method, "path": path, "json": json, "params": params} ) @@ -102,15 +103,15 @@ async def _replace_async( # NOSONAR(S7503) — must be async def to match the a *, method: str, path: str, - json: Optional[Dict[str, Any]] = None, - params: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: + json: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + ) -> dict[str, Any]: return self._replace_sync( method=method, path=path, json=json, params=params ) @property - def last_body(self) -> Optional[Dict[str, Any]]: + def last_body(self) -> dict[str, Any] | None: if not self.calls: return None return self.calls[-1]["json"] @@ -119,7 +120,7 @@ def last_body(self) -> Optional[Dict[str, Any]]: @pytest.fixture def http_client_with_capture( monkeypatch: pytest.MonkeyPatch, -) -> Tuple[SlayerClient, _CapturedRequests]: +) -> tuple[SlayerClient, _CapturedRequests]: """A remote-mode ``SlayerClient`` whose request-wrappers are stubbed. Returns ``(client, capture)`` where ``capture.last_body`` exposes the @@ -282,7 +283,7 @@ class TestHttpBodyShape: def test_list_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture queries = [ @@ -310,7 +311,7 @@ def test_list_body_shape( def test_list_with_slayerquery_items( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """Mixed list items (``SlayerQuery`` + dict) each serialise to dict.""" client, cap = http_client_with_capture @@ -319,7 +320,7 @@ def test_list_with_slayerquery_items( source_model="orders", measures=[{"formula": "amount:sum"}], ) - items: List[Any] = [q, {"source_model": "a"}] + items: list[Any] = [q, {"source_model": "a"}] client.query_sync(items) body = cap.last_body assert body is not None @@ -339,7 +340,7 @@ def test_list_with_slayerquery_items( def test_str_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture client.query_sync("rev_by_region") @@ -349,7 +350,7 @@ def test_str_body_shape( def test_dict_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """Dict input is round-tripped through ``SlayerQuery`` so the server sees the JSON-mode dump (string-shorthand normalised, @@ -373,7 +374,7 @@ def test_dict_body_shape( def test_dict_normalizes_string_shorthand( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """String-shorthand measures (e.g. ``"amount:sum"``) become dict-form (``{"formula": "amount:sum"}``) before reaching the @@ -394,7 +395,7 @@ def test_dict_normalizes_string_shorthand( def test_slayerquery_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture q = SlayerQuery( @@ -408,7 +409,7 @@ def test_slayerquery_body_shape( def test_dry_run_explain_appended_list( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture items = [ @@ -428,7 +429,7 @@ def test_dry_run_explain_appended_list( def test_dry_run_explain_appended_str( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture client.query_sync("m", dry_run=True, explain=True) @@ -440,7 +441,7 @@ def test_dry_run_explain_appended_str( def test_dry_run_explain_appended_dict( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture client.query_sync( @@ -454,7 +455,7 @@ def test_dry_run_explain_appended_dict( def test_flags_omitted_when_false( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """``dry_run=False`` / ``explain=False`` (defaults) → keys not present.""" client, cap = http_client_with_capture @@ -468,7 +469,7 @@ def test_flags_omitted_when_false( def test_does_not_mutate_caller_dict( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, _cap = http_client_with_capture payload = { @@ -483,10 +484,10 @@ def test_does_not_mutate_caller_dict( def test_does_not_mutate_caller_list( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture - items: List[Dict[str, Any]] = [ + items: list[dict[str, Any]] = [ {"name": "a", "source_model": "orders"}, {"source_model": "a"}, ] @@ -506,7 +507,7 @@ def test_does_not_mutate_caller_list( def test_rejects_invalid_input_top_level( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture with pytest.raises(TypeError, match="SlayerQuery"): @@ -515,7 +516,7 @@ def test_rejects_invalid_input_top_level( def test_rejects_invalid_list_item( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture with pytest.raises(TypeError, match=r"query\[1\]"): @@ -528,7 +529,7 @@ def test_rejects_invalid_list_item( def test_mappingproxy_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """``MappingProxyType`` is a ``Mapping`` but not a ``dict`` — the helper must honour the declared ``Mapping[str, Any]`` contract @@ -544,7 +545,7 @@ def test_mappingproxy_body_shape( def test_tuple_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """``tuple`` is a ``Sequence`` but not a ``list`` — honoured.""" client, cap = http_client_with_capture @@ -567,7 +568,7 @@ def test_tuple_body_shape( def test_dict_with_variables_preserved( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """A dict input carrying top-level ``variables`` reaches the server verbatim. DEV-1437 doesn't add a ``variables=`` kwarg (see DEV-1438) @@ -588,7 +589,7 @@ def test_dict_with_variables_preserved( def test_sql_sync_list_input( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """``sql_sync`` delegates to ``query_sync(dry_run=True)`` — list input must reach the transport with ``dry_run`` set.""" @@ -612,7 +613,7 @@ def test_sql_sync_list_input( def test_sql_sync_str_input( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture sql = client.sql_sync("rev_by_region") @@ -621,7 +622,7 @@ def test_sql_sync_str_input( def test_explain_sync_list_input( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """``explain_sync`` delegates to ``query_sync(explain=True)`` — list input must reach the transport with ``explain`` set.""" @@ -645,7 +646,7 @@ def test_explain_sync_list_input( def test_explain_sync_str_input( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture client.explain_sync("rev_by_region") @@ -666,9 +667,9 @@ def _replace_sync_with_rows( *, method: str, path: str, - json: Optional[Dict[str, Any]] = None, - params: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: + json: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + ) -> dict[str, Any]: capture.calls.append( {"method": method, "path": path, "json": json, "params": params} ) @@ -702,7 +703,7 @@ def _replace_sync_with_rows( async def test_async_query_list_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: """Async ``query`` mirrors sync ``query_sync`` body shape.""" client, cap = http_client_with_capture @@ -723,7 +724,7 @@ async def test_async_query_list_body_shape( async def test_async_query_str_body_shape( self, - http_client_with_capture: Tuple[SlayerClient, _CapturedRequests], + http_client_with_capture: tuple[SlayerClient, _CapturedRequests], ) -> None: client, cap = http_client_with_capture await client.query("rev_by_region") diff --git a/tests/test_column_dependency.py b/tests/test_column_dependency.py index ec8cc0ec..8f234d07 100644 --- a/tests/test_column_dependency.py +++ b/tests/test_column_dependency.py @@ -10,7 +10,6 @@ create/edit, and the migration write-back (the migration path passes ``_validate=False`` so legacy cyclic data remains loadable). """ -from typing import Tuple import pytest @@ -53,7 +52,7 @@ def _abc_chain( a_foo_sql: str, c_x_sql: str, c_joins_back_to_a: bool = False, -) -> Tuple[SlayerModel, SlayerModel, SlayerModel]: +) -> tuple[SlayerModel, SlayerModel, SlayerModel]: """The canonical A → B → C three-model scaffold used by the strict alias-walk tests. ``foo`` on A and ``x`` on C are derived; their expressions are the per-test variable. diff --git a/tests/test_edit_datasource_refresh.py b/tests/test_edit_datasource_refresh.py index 39f0271f..a6eb53fb 100644 --- a/tests/test_edit_datasource_refresh.py +++ b/tests/test_edit_datasource_refresh.py @@ -7,7 +7,7 @@ import shutil import tempfile -from typing import Generator +from collections.abc import Generator import pytest diff --git a/tests/test_embedding_retriever.py b/tests/test_embedding_retriever.py index 5ba48a3b..42380c2b 100644 --- a/tests/test_embedding_retriever.py +++ b/tests/test_embedding_retriever.py @@ -19,7 +19,7 @@ from __future__ import annotations import tempfile -from typing import List, Optional, cast +from typing import cast import pytest @@ -46,14 +46,14 @@ def stub_available(monkeypatch: pytest.MonkeyPatch) -> None: class _RecordingEmbedBatch: def __init__(self) -> None: - self.calls: List[List[str]] = [] + self.calls: list[list[str]] = [] self.override_none: set = set() async def __call__( - self, texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + self, texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: self.calls.append(list(texts)) - out: List[Optional[List[float]]] = [] + out: list[list[float] | None] = [] for global_idx, text in enumerate(texts): if (len(self.calls) - 1, global_idx) in self.override_none: out.append(None) @@ -210,9 +210,9 @@ async def test_refresh_model_subtree_per_entry_failure_warns( monkeypatch: pytest.MonkeyPatch, ) -> None: async def partial_failure( # NOSONAR(S7503) - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: - out: List[Optional[List[float]]] = [] + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: + out: list[list[float] | None] = [] for i, _ in enumerate(texts): out.append(None if i == 0 else [0.1, 0.2, 0.3]) return out @@ -330,8 +330,8 @@ async def test_apply_pending_persists_partial_batch_on_some_embed_failures( monkeypatch: pytest.MonkeyPatch, ) -> None: async def partial_failure( # NOSONAR(S7503) - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: return [None] + [[0.1, 0.2, 0.3]] * (len(texts) - 1) monkeypatch.setattr( diff --git a/tests/test_embeddings_client.py b/tests/test_embeddings_client.py index 162abbc0..61746f43 100644 --- a/tests/test_embeddings_client.py +++ b/tests/test_embeddings_client.py @@ -8,7 +8,7 @@ from __future__ import annotations import asyncio -from typing import Any, List, Optional, Tuple +from typing import Any import pytest @@ -68,10 +68,10 @@ async def test_embed_batch_calls_litellm_with_resolved_model( captured: dict = {} class _FakeResponse: - def __init__(self, data: List[dict]) -> None: + def __init__(self, data: list[dict]) -> None: self.data = data - async def fake_aembedding(*, model: str, input: List[str]) -> _FakeResponse: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> _FakeResponse: # NOSONAR(S7503) — stub matches litellm.aembedding async signature captured["model"] = model captured["input"] = list(input) return _FakeResponse( @@ -129,8 +129,8 @@ async def test_embed_query_caches_repeated_calls( call_count = {"n": 0} async def fake_embed_batch( # NOSONAR(S7503) — stub matches embed_batch async signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: call_count["n"] += 1 return [[float(call_count["n"])] * 3 for _ in texts] @@ -311,9 +311,9 @@ def test_truncate_text_for_model_get_max_tokens_prefixed_then_bare( falling back to 8192.""" pytest.importorskip("litellm") import litellm - calls: List[str] = [] + calls: list[str] = [] - def fake_get_max(model: str) -> Optional[int]: + def fake_get_max(model: str) -> int | None: calls.append(model) if model == "openai/text-embedding-3-small": return None @@ -470,7 +470,7 @@ class _Resp: def __init__(self, n: int) -> None: self.data = [{"embedding": [float(i)] * 3} for i in range(n)] - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature return _Resp(len(input)) monkeypatch.setattr(litellm, "aembedding", fake_aembedding) @@ -490,9 +490,9 @@ def test_truncate_text_for_model_get_max_tokens_raises_then_bare_used( pytest.importorskip("litellm") tiktoken = pytest.importorskip("tiktoken") import litellm - calls: List[str] = [] + calls: list[str] = [] - def fake_get_max(model: str) -> Optional[int]: + def fake_get_max(model: str) -> int | None: calls.append(model) if model == "openai/text-embedding-3-small": raise RuntimeError("transient") @@ -545,7 +545,7 @@ def test_truncate_text_for_model_failure_not_cached( import litellm call = {"n": 0} - def fake_get_max(_m: str) -> Optional[int]: + def fake_get_max(_m: str) -> int | None: call["n"] += 1 if call["n"] == 1: raise RuntimeError("transient") @@ -605,7 +605,7 @@ def _install_per_input_aembedding_stub( monkeypatch: pytest.MonkeyPatch, *, on_single: Any, -) -> Tuple[List[List[str]], List[str]]: +) -> tuple[list[list[str]], list[str]]: """Wire up a fake ``litellm.aembedding`` for per-input-retry tests. The fake records every (input, model) the implementation calls it @@ -618,10 +618,10 @@ def _install_per_input_aembedding_stub( """ litellm = pytest.importorskip("litellm") bad = _make_bad_request_error() - seen_inputs: List[List[str]] = [] - seen_models: List[str] = [] + seen_inputs: list[list[str]] = [] + seen_models: list[str] = [] - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature seen_inputs.append(list(input)) seen_models.append(model) if len(input) > 1: @@ -645,13 +645,13 @@ async def test_embed_batch_truncates_over_cap_inputs_happy_path( monkeypatch.setattr(litellm.utils, "get_max_tokens", lambda _m: 300) long_text = "word " * 2000 # over the 44-token budget - captured_inputs: List[List[str]] = [] + captured_inputs: list[list[str]] = [] class _Resp: def __init__(self, n: int) -> None: self.data = [{"embedding": [float(i)] * 3} for i in range(n)] - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature captured_inputs.append(list(input)) return _Resp(len(input)) @@ -754,10 +754,10 @@ async def test_embed_batch_batch_non_bad_request_returns_all_none( monkeypatch.setattr(litellm.utils, "get_max_tokens", lambda _m: 300) call_count = {"n": 0} - captured_inputs: List[List[str]] = [] - captured_models: List[str] = [] + captured_inputs: list[list[str]] = [] + captured_models: list[str] = [] - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature call_count["n"] += 1 captured_inputs.append(list(input)) captured_models.append(model) @@ -817,11 +817,11 @@ async def test_embed_batch_no_truncation_when_unavailable( real_truncate = embedding_client.truncate_text_for_model - def spy_truncate(text: str, *, model: Optional[str] = None) -> str: + def spy_truncate(text: str, *, model: str | None = None) -> str: truncate_calls["n"] += 1 return real_truncate(text, model=model) - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature aembedding_calls["n"] += 1 return None @@ -858,7 +858,7 @@ class _Resp: def __init__(self) -> None: self.data = [{"embedding": [1.0, 2.0, 3.0]}] - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature captured["input"] = list(input) return _Resp() @@ -890,7 +890,7 @@ async def test_embed_query_returns_none_when_per_input_retry_fails( bad = _make_bad_request_error() call_count = {"n": 0} - async def fake_aembedding(*, model: str, input: List[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature + async def fake_aembedding(*, model: str, input: list[str]) -> Any: # NOSONAR(S7503) — stub matches litellm.aembedding async signature call_count["n"] += 1 raise bad diff --git a/tests/test_embeddings_storage.py b/tests/test_embeddings_storage.py index 3eb44d5d..464a6f55 100644 --- a/tests/test_embeddings_storage.py +++ b/tests/test_embeddings_storage.py @@ -11,7 +11,7 @@ import os import tempfile from datetime import datetime, timezone -from typing import Iterator, List +from collections.abc import Iterator import pytest @@ -44,7 +44,7 @@ def _embed( model: str = "openai/test-embedding", kind: EntityKind = "memory", text_hash: str = "h0", - vector: List[float] | None = None, + vector: list[float] | None = None, ) -> Embedding: return Embedding( canonical_id=canonical_id, @@ -201,7 +201,7 @@ async def test_delete_model_cascades_subtree_embeddings( ) await storage.save_model(model) - seeds: List[tuple[str, EntityKind]] = [ + seeds: list[tuple[str, EntityKind]] = [ ("dsx.orders", "model"), ("dsx.orders.id", "column"), ("dsx.orders.rev", "measure"), @@ -232,7 +232,7 @@ async def test_delete_datasource_cascades_descendants( name="dsx", type="postgres", host="h", database="d", ) await storage.save_datasource(ds) - seeds: List[tuple[str, EntityKind]] = [ + seeds: list[tuple[str, EntityKind]] = [ ("dsx", "datasource"), ("dsx.orders", "model"), ("dsx.orders.id", "column"), diff --git a/tests/test_engine_profiling.py b/tests/test_engine_profiling.py index 2a1673c5..65fb9384 100644 --- a/tests/test_engine_profiling.py +++ b/tests/test_engine_profiling.py @@ -22,7 +22,6 @@ import asyncio import sqlite3 import tempfile -from typing import Optional import pytest @@ -781,7 +780,7 @@ async def test_refresh_continues_after_per_column_failure(sqlite_setup, monkeypa call_count = {"n": 0} real_profile_column = profile_column - async def boom_then_ok(*, model, column, engine) -> Optional[ColumnSample]: + async def boom_then_ok(*, model, column, engine) -> ColumnSample | None: call_count["n"] += 1 if column.name == "amount": raise RuntimeError("simulated profile failure") diff --git a/tests/test_entity_description_plumbing.py b/tests/test_entity_description_plumbing.py index 3fa1ceda..9352e975 100644 --- a/tests/test_entity_description_plumbing.py +++ b/tests/test_entity_description_plumbing.py @@ -16,7 +16,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_entity_resolution.py b/tests/test_entity_resolution.py index 4f793f74..774394db 100644 --- a/tests/test_entity_resolution.py +++ b/tests/test_entity_resolution.py @@ -16,7 +16,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_entry_points_namespacing.py b/tests/test_entry_points_namespacing.py index 16ad694b..13061e81 100644 --- a/tests/test_entry_points_namespacing.py +++ b/tests/test_entry_points_namespacing.py @@ -20,7 +20,7 @@ """ import tempfile -from typing import Any, Optional +from typing import Any import pytest from fastapi.testclient import TestClient @@ -67,7 +67,7 @@ def _ds(name: str) -> DatasourceConfig: return DatasourceConfig(name=name, type="postgres", host="h") -async def _call_mcp(server, *, name: str, arguments: Optional[dict[str, Any]] = None) -> str: +async def _call_mcp(server, *, name: str, arguments: dict[str, Any] | None = None) -> str: """Invoke an MCP tool and return its text result.""" blocks, _ = await server.call_tool(name=name, arguments=arguments or {}) return blocks[0].text diff --git a/tests/test_force_clean.py b/tests/test_force_clean.py index ace84bd1..a2e44a9e 100644 --- a/tests/test_force_clean.py +++ b/tests/test_force_clean.py @@ -10,7 +10,6 @@ import sys import tempfile from pathlib import Path -from typing import List from unittest.mock import patch import pytest @@ -102,7 +101,7 @@ async def _setup(workspace: Path) -> tuple[SlayerQueryEngine, str]: class TestApplyDriftDeletes: async def test_applies_edit_model_delete(self, workspace: Path) -> None: engine, _ = await _setup(workspace) - deletes: List[ToDeleteEntry] = [ + deletes: list[ToDeleteEntry] = [ EditModelDelete( model_name="orders", data_source="ds", @@ -122,7 +121,7 @@ async def test_applies_edit_model_delete(self, workspace: Path) -> None: async def test_applies_whole_model_delete(self, workspace: Path) -> None: engine, _ = await _setup(workspace) - deletes: List[ToDeleteEntry] = [ + deletes: list[ToDeleteEntry] = [ WholeModelDelete( model_name="orders", data_source="ds", diff --git a/tests/test_idempotent_ingestion.py b/tests/test_idempotent_ingestion.py index 70602dab..d8b277c8 100644 --- a/tests/test_idempotent_ingestion.py +++ b/tests/test_idempotent_ingestion.py @@ -16,7 +16,8 @@ import sqlite3 import tempfile from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional +from typing import Any +from collections.abc import Iterable import pytest @@ -384,15 +385,15 @@ class TestMemoryEmbeddingRefresh: datasource doc. Failures attribute to the offending memory by id.""" @staticmethod - def _enable_channel(monkeypatch: pytest.MonkeyPatch) -> List[List[str]]: + def _enable_channel(monkeypatch: pytest.MonkeyPatch) -> list[list[str]]: """Override the conftest autouse fixture for this test. Returns a list that captures every ``embed_batch`` call's text payload.""" monkeypatch.setattr(embedding_client, "is_available", lambda: True) - calls: List[List[str]] = [] + calls: list[list[str]] = [] async def fake_embed_batch( # NOSONAR(S7503) — must be `async def` to match the patched embed_batch signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: calls.append(list(texts)) return [[0.1, 0.2, 0.3] for _ in texts] @@ -557,11 +558,11 @@ async def test_extra_not_installed_silent_noop_on_memory_path( # clarity. assert embedding_client.is_available() is False - called: List[List[str]] = [] + called: list[list[str]] = [] async def should_not_be_called( # NOSONAR(S7503) — must be `async def` to match the patched embed_batch signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: called.append(list(texts)) return [None for _ in texts] @@ -694,14 +695,14 @@ def _create_probe_workspace_db( async def _persist_int_model( storage, ds_name: str, table: str, column: str, *, - column_format: Optional[NumberFormat] = None, - extra_meta: Optional[dict] = None, + column_format: NumberFormat | None = None, + extra_meta: dict | None = None, description: str = "", - label: Optional[str] = None, + label: str | None = None, ) -> None: """Persist a model with the target column hard-coded as type=INT, simulating a pre-DEV-1538 ingest.""" - col_kwargs: Dict[str, Any] = { + col_kwargs: dict[str, Any] = { "name": column, "sql": column, "type": DataType.INT, diff --git a/tests/test_ingestion_memory_dangling_cleanup.py b/tests/test_ingestion_memory_dangling_cleanup.py index 6f446e89..276bd25e 100644 --- a/tests/test_ingestion_memory_dangling_cleanup.py +++ b/tests/test_ingestion_memory_dangling_cleanup.py @@ -11,7 +11,7 @@ import os import sqlite3 import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest from sqlalchemy.exc import OperationalError as SAOperationalError diff --git a/tests/test_inspect_model_compact.py b/tests/test_inspect_model_compact.py index 75063ca2..9a863b06 100644 --- a/tests/test_inspect_model_compact.py +++ b/tests/test_inspect_model_compact.py @@ -11,7 +11,8 @@ import json import shutil import tempfile -from typing import Any, Generator +from typing import Any +from collections.abc import Generator import pytest diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 074dacf4..163eee3d 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -4,7 +4,8 @@ import os import shutil import tempfile -from typing import Any, Generator, Optional +from typing import Any +from collections.abc import Generator import pytest @@ -83,7 +84,7 @@ def mcp_server(_shared_mcp_server, storage: YAMLStorage): return _shared_mcp_server -async def _call(mcp_server, *, name: str, arguments: Optional[dict[str, Any]] = None) -> str: +async def _call(mcp_server, *, name: str, arguments: dict[str, Any] | None = None) -> str: """Call an MCP tool and return the text result.""" content_blocks, result_dict = await mcp_server.call_tool(name=name, arguments=arguments or {}) return content_blocks[0].text diff --git a/tests/test_memories_cli.py b/tests/test_memories_cli.py index 70cb54eb..bdceca05 100644 --- a/tests/test_memories_cli.py +++ b/tests/test_memories_cli.py @@ -19,7 +19,6 @@ import sys import tempfile from types import SimpleNamespace -from typing import Optional import pytest @@ -75,10 +74,10 @@ def _args( *, storage_path: str, memory_command: str, - learning: Optional[str] = None, - entities: Optional[str] = None, - query: Optional[str] = None, - id: Optional[int] = None, # noqa: A002 — argparse arg + learning: str | None = None, + entities: str | None = None, + query: str | None = None, + id: int | None = None, # noqa: A002 — argparse arg ) -> SimpleNamespace: return SimpleNamespace( storage=storage_path, diff --git a/tests/test_memories_client.py b/tests/test_memories_client.py index e0a66363..e90b0a27 100644 --- a/tests/test_memories_client.py +++ b/tests/test_memories_client.py @@ -9,7 +9,7 @@ """ import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_memories_inspect.py b/tests/test_memories_inspect.py index 4e57231b..cd0dcaeb 100644 --- a/tests/test_memories_inspect.py +++ b/tests/test_memories_inspect.py @@ -12,7 +12,8 @@ import shutil import sqlite3 import tempfile -from typing import Any, Generator, Optional +from typing import Any +from collections.abc import Generator import pytest @@ -120,7 +121,7 @@ async def _call( mcp_server, *, name: str, - arguments: Optional[dict[str, Any]] = None, + arguments: dict[str, Any] | None = None, ) -> str: content_blocks, _ = await mcp_server.call_tool( name=name, arguments=arguments or {} diff --git a/tests/test_memories_mcp.py b/tests/test_memories_mcp.py index bd36136d..de80a6d2 100644 --- a/tests/test_memories_mcp.py +++ b/tests/test_memories_mcp.py @@ -19,7 +19,8 @@ import shutil import sqlite3 import tempfile -from typing import Any, Generator, Optional +from typing import Any +from collections.abc import Generator import pytest @@ -162,7 +163,7 @@ async def _call( mcp_server, *, name: str, - arguments: Optional[dict[str, Any]] = None, + arguments: dict[str, Any] | None = None, ) -> str: content_blocks, _ = await mcp_server.call_tool( name=name, arguments=arguments or {} @@ -170,7 +171,7 @@ async def _call( return content_blocks[0].text -def _try_parse_json(text: str) -> Optional[dict]: +def _try_parse_json(text: str) -> dict | None: try: return json.loads(text) except json.JSONDecodeError: diff --git a/tests/test_memories_rest.py b/tests/test_memories_rest.py index 432342ce..4cf28aae 100644 --- a/tests/test_memories_rest.py +++ b/tests/test_memories_rest.py @@ -16,7 +16,7 @@ import shutil import sqlite3 import tempfile -from typing import Generator +from collections.abc import Generator import pytest from fastapi.testclient import TestClient diff --git a/tests/test_memories_storage.py b/tests/test_memories_storage.py index 1c9e70f6..dfb1e409 100644 --- a/tests/test_memories_storage.py +++ b/tests/test_memories_storage.py @@ -18,7 +18,7 @@ import asyncio import os import tempfile -from typing import Iterator +from collections.abc import Iterator import pytest diff --git a/tests/test_memory_resolver_memory_branch.py b/tests/test_memory_resolver_memory_branch.py index b6897d57..a39d31d8 100644 --- a/tests/test_memory_resolver_memory_branch.py +++ b/tests/test_memory_resolver_memory_branch.py @@ -9,7 +9,7 @@ import os import tempfile -from typing import Iterator +from collections.abc import Iterator import pytest diff --git a/tests/test_memory_string_ids.py b/tests/test_memory_string_ids.py index 6d211024..e816fc2b 100644 --- a/tests/test_memory_string_ids.py +++ b/tests/test_memory_string_ids.py @@ -10,7 +10,7 @@ import os import tempfile -from typing import Iterator +from collections.abc import Iterator import pytest diff --git a/tests/test_models_summary_compact.py b/tests/test_models_summary_compact.py index 6565f80d..69518cf4 100644 --- a/tests/test_models_summary_compact.py +++ b/tests/test_models_summary_compact.py @@ -20,7 +20,8 @@ import json import shutil import tempfile -from typing import Any, Generator +from typing import Any +from collections.abc import Generator import pytest diff --git a/tests/test_nested_dag_cross_stage_refs.py b/tests/test_nested_dag_cross_stage_refs.py index a2d92ac9..99b940ab 100644 --- a/tests/test_nested_dag_cross_stage_refs.py +++ b/tests/test_nested_dag_cross_stage_refs.py @@ -17,7 +17,6 @@ import sqlite3 import tempfile -from typing import List, Tuple import pytest import sqlglot @@ -105,7 +104,7 @@ def _orders_model() -> SlayerModel: ) -async def _engine_with_join_chain() -> Tuple[SlayerQueryEngine, tempfile.TemporaryDirectory]: +async def _engine_with_join_chain() -> tuple[SlayerQueryEngine, tempfile.TemporaryDirectory]: """Build a YAMLStorage with orders → customers → regions → countries.""" tmp = tempfile.TemporaryDirectory() storage = YAMLStorage(base_dir=tmp.name) @@ -181,7 +180,7 @@ def _outermost_select(sql: str, dialect: str = "sqlite") -> exp.Select: return tree -def _outer_column_refs(sql: str, dialect: str = "sqlite") -> List[exp.Column]: +def _outer_column_refs(sql: str, dialect: str = "sqlite") -> list[exp.Column]: """All `exp.Column` nodes whose nearest ancestor SELECT is the OUTERMOST SELECT of `sql` — i.e. columns directly in the outer projection / GROUP BY / ORDER BY / HAVING / WHERE. @@ -189,7 +188,7 @@ def _outer_column_refs(sql: str, dialect: str = "sqlite") -> List[exp.Column]: Does NOT include columns inside subqueries / CTEs / inner SELECTs. """ outer = _outermost_select(sql, dialect=dialect) - nodes: List[exp.Column] = [] + nodes: list[exp.Column] = [] for col in outer.find_all(exp.Column): if col.parent_select is outer: nodes.append(col) diff --git a/tests/test_projection_trim.py b/tests/test_projection_trim.py index 55ad6cd3..5ca42b87 100644 --- a/tests/test_projection_trim.py +++ b/tests/test_projection_trim.py @@ -13,7 +13,7 @@ import importlib import inspect import re -from typing import Any, List +from typing import Any import pytest import sqlglot @@ -57,7 +57,7 @@ def _norm(s: str) -> str: return " ".join(s.split()) -def _outer_select_columns(sql: str, *, dialect: str = "postgres") -> List[str]: +def _outer_select_columns(sql: str, *, dialect: str = "postgres") -> list[str]: """Return the list of alias names projected by the OUTERMOST SELECT. For ``SELECT a, b FROM (...)`` the outer projection is ``[a, b]``. @@ -68,14 +68,14 @@ def _outer_select_columns(sql: str, *, dialect: str = "postgres") -> List[str]: parsed = sqlglot.parse_one(sql, dialect=dialect) if not isinstance(parsed, exp.Select): # pragma: no cover — defensive return [] - out: List[str] = [] + out: list[str] = [] for proj in parsed.expressions: # ``alias_or_name`` returns the alias if present, else the bare name. out.append(proj.alias_or_name) return out -def _all_aliases_in_sql(sql: str) -> List[str]: +def _all_aliases_in_sql(sql: str) -> list[str]: """Return every alias-name that appears as a quoted identifier in the SQL. Useful for "alias X appears in some CTE" assertions without committing @@ -84,7 +84,7 @@ def _all_aliases_in_sql(sql: str) -> List[str]: return re.findall(r'"([^"]+)"', sql) -def _outer_order_by_references(sql: str, *, dialect: str = "postgres") -> List[str]: +def _outer_order_by_references(sql: str, *, dialect: str = "postgres") -> list[str]: """Return identifier names referenced by the outermost ORDER BY clause.""" parsed = sqlglot.parse_one(sql, dialect=dialect) if not isinstance(parsed, exp.Select): # pragma: no cover @@ -92,7 +92,7 @@ def _outer_order_by_references(sql: str, *, dialect: str = "postgres") -> List[s order = parsed.args.get("order") if order is None: return [] - refs: List[str] = [] + refs: list[str] = [] for ordered in order.expressions: col = ordered.this if isinstance(col, exp.Column): @@ -1431,7 +1431,7 @@ async def test_get_column_types_uses_outer_render_mode( # Spy on the generate calls to capture the render_mode used by # get_column_types. We replace the bound method on the class with # a wrapper that records the kwargs. - captured: List[dict] = [] + captured: list[dict] = [] original_generate = SQLGenerator.generate def _wrapper(self, *args, **kwargs): diff --git a/tests/test_rest_memory_string_ids.py b/tests/test_rest_memory_string_ids.py index cd1f62c5..4caaf108 100644 --- a/tests/test_rest_memory_string_ids.py +++ b/tests/test_rest_memory_string_ids.py @@ -9,7 +9,7 @@ import os import tempfile -from typing import Iterator +from collections.abc import Iterator import pytest from fastapi.testclient import TestClient diff --git a/tests/test_retriever_abc.py b/tests/test_retriever_abc.py index 7e7d6c1f..1635fcdc 100644 --- a/tests/test_retriever_abc.py +++ b/tests/test_retriever_abc.py @@ -9,7 +9,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_save_memory_description_plumbing.py b/tests/test_save_memory_description_plumbing.py index ed40ac3b..4ea6644f 100644 --- a/tests/test_save_memory_description_plumbing.py +++ b/tests/test_save_memory_description_plumbing.py @@ -8,7 +8,8 @@ import json import shutil import tempfile -from typing import Any, AsyncIterator, Generator +from typing import Any +from collections.abc import AsyncIterator, Generator import pytest import pytest_asyncio diff --git a/tests/test_save_memory_id_kwarg.py b/tests/test_save_memory_id_kwarg.py index f3300cce..0cfb3431 100644 --- a/tests/test_save_memory_id_kwarg.py +++ b/tests/test_save_memory_id_kwarg.py @@ -10,7 +10,7 @@ import asyncio import os import tempfile -from typing import Iterator +from collections.abc import Iterator import pytest diff --git a/tests/test_save_memory_linked_entities_memory_refs.py b/tests/test_save_memory_linked_entities_memory_refs.py index 7a5680e3..cf262662 100644 --- a/tests/test_save_memory_linked_entities_memory_refs.py +++ b/tests/test_save_memory_linked_entities_memory_refs.py @@ -9,7 +9,7 @@ import os import tempfile -from typing import Iterator +from collections.abc import Iterator import pytest diff --git a/tests/test_search_compact.py b/tests/test_search_compact.py index a2e250dd..b7d8e012 100644 --- a/tests/test_search_compact.py +++ b/tests/test_search_compact.py @@ -25,7 +25,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_search_compact_plumbing.py b/tests/test_search_compact_plumbing.py index 76a99ff3..a8eba023 100644 --- a/tests/test_search_compact_plumbing.py +++ b/tests/test_search_compact_plumbing.py @@ -13,7 +13,8 @@ import sys import tempfile from types import SimpleNamespace -from typing import Any, AsyncIterator, Generator +from typing import Any +from collections.abc import AsyncIterator, Generator import pytest import pytest_asyncio diff --git a/tests/test_search_datasource_filter.py b/tests/test_search_datasource_filter.py index ebe00998..2275b5ea 100644 --- a/tests/test_search_datasource_filter.py +++ b/tests/test_search_datasource_filter.py @@ -19,7 +19,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_search_description_indexed.py b/tests/test_search_description_indexed.py index 29547ae5..e6a39d40 100644 --- a/tests/test_search_description_indexed.py +++ b/tests/test_search_description_indexed.py @@ -7,7 +7,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_search_graph.py b/tests/test_search_graph.py index a6761da8..ce83af46 100644 --- a/tests/test_search_graph.py +++ b/tests/test_search_graph.py @@ -14,7 +14,7 @@ import asyncio import os import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator from unittest.mock import patch import pytest diff --git a/tests/test_search_index.py b/tests/test_search_index.py index 292e20d2..8288ca63 100644 --- a/tests/test_search_index.py +++ b/tests/test_search_index.py @@ -17,7 +17,6 @@ from __future__ import annotations -from typing import List import pytest @@ -27,7 +26,7 @@ from slayer.search.index import build_in_memory_index, search_index -def _make_models() -> List[SlayerModel]: +def _make_models() -> list[SlayerModel]: return [ SlayerModel( name="orders", @@ -61,7 +60,7 @@ def _make_models() -> List[SlayerModel]: ] -def _make_memories() -> List[Memory]: +def _make_memories() -> list[Memory]: return [ Memory( id=1, diff --git a/tests/test_search_invariance.py b/tests/test_search_invariance.py index da90c019..8288c7ce 100644 --- a/tests/test_search_invariance.py +++ b/tests/test_search_invariance.py @@ -16,7 +16,7 @@ import hashlib import tempfile -from typing import AsyncIterator, List, Optional +from collections.abc import AsyncIterator import pytest import pytest_asyncio @@ -84,7 +84,7 @@ ] -def _make_models() -> List[SlayerModel]: +def _make_models() -> list[SlayerModel]: return [ SlayerModel( name="orders", @@ -148,7 +148,7 @@ def _make_models() -> List[SlayerModel]: ] -def _entities_for_topic(topic: str) -> List[str]: +def _entities_for_topic(topic: str) -> list[str]: """Pick canonical entity tags for a learning-topic string. Pulled out of ``_seed_invariance_corpus`` so each branch stays separate from the seeding loop's control flow.""" @@ -372,8 +372,8 @@ async def storage_with_embeddings( # vector so ranks vary across docs but are reproducible across # interpreter runs (Python's built-in ``hash`` is randomised # per process, so use sha256 here). - def _vec(text: str) -> List[float]: - out: List[float] = [] + def _vec(text: str) -> list[float]: + out: list[float] = [] for i in range(8): digest = hashlib.sha256( f"{text}|{i}".encode("utf-8"), @@ -383,13 +383,13 @@ def _vec(text: str) -> List[float]: return out async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch async signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: return [_vec(t) for t in texts] async def stub_embed_query( # NOSONAR(S7503) — stub matches embed_query async signature - text: str, *, model: Optional[str] = None, - ) -> List[float]: + text: str, *, model: str | None = None, + ) -> list[float]: return _vec(text) monkeypatch.setattr( diff --git a/tests/test_search_named_entity_surfacing.py b/tests/test_search_named_entity_surfacing.py index e604fdca..eefa7df8 100644 --- a/tests/test_search_named_entity_surfacing.py +++ b/tests/test_search_named_entity_surfacing.py @@ -19,7 +19,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator, List +from collections.abc import AsyncIterator import pytest @@ -41,7 +41,7 @@ # --------------------------------------------------------------------------- -def _make_corpus_models() -> List[SlayerModel]: +def _make_corpus_models() -> list[SlayerModel]: return [ SlayerModel( name="orders", diff --git a/tests/test_search_render_unified.py b/tests/test_search_render_unified.py index 138d91f2..a72308e6 100644 --- a/tests/test_search_render_unified.py +++ b/tests/test_search_render_unified.py @@ -23,7 +23,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator, List, Optional +from collections.abc import AsyncIterator import pytest @@ -182,11 +182,11 @@ async def test_embedding_refresh_model_subtree_uses_unified_helper( dispatch will fail.""" from slayer.search.render import collect_model_entity_pairs - captured: List[str] = [] + captured: list[str] = [] async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: # We only need to capture the rendered texts; the embedding # value itself is irrelevant for this parity assertion. captured.extend(texts) @@ -247,11 +247,11 @@ async def test_render_datasource_pair_used_by_index_and_embedding_paths( assert pair_ws.text == corpus_ws_text # Embedding-side: stub embed_batch and inspect the captured text. - captured: List[str] = [] + captured: list[str] = [] async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: captured.extend(texts) return [[0.0, 1.0]] * len(texts) @@ -313,7 +313,7 @@ def test_index_corpus_path_actually_calls_collect_model_entity_pairs( canonical_id="SENTINEL.x", kind="model", text="SENTINEL_TEXT", ) - def stub(*, model: SlayerModel) -> List[RenderedEntity]: # noqa: ARG001 + def stub(*, model: SlayerModel) -> list[RenderedEntity]: # noqa: ARG001 return [sentinel] monkeypatch.setattr(index_mod, "collect_model_entity_pairs", stub) @@ -345,16 +345,16 @@ async def test_embedding_path_actually_calls_collect_model_entity_pairs( canonical_id="SENTINEL.x", kind="model", text=sentinel_text, ) - def stub(*, model: SlayerModel) -> List[RenderedEntity]: # noqa: ARG001 + def stub(*, model: SlayerModel) -> list[RenderedEntity]: # noqa: ARG001 return [sentinel] monkeypatch.setattr(emb_mod, "collect_model_entity_pairs", stub) - captured: List[str] = [] + captured: list[str] = [] async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch signature - texts: List[str], *, model: Optional[str] = None, # noqa: ARG001 - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, # noqa: ARG001 + ) -> list[list[float] | None]: captured.extend(texts) return [[0.0, 1.0]] * len(texts) @@ -382,8 +382,8 @@ def test_index_corpus_path_actually_calls_render_datasource_pair( def stub( *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, + models: list[SlayerModel], + description: str | None = None, ) -> RenderedEntity: # noqa: ARG001 return sentinel @@ -415,18 +415,18 @@ async def test_embedding_path_actually_calls_render_datasource_pair( def stub( *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, + models: list[SlayerModel], + description: str | None = None, ) -> RenderedEntity: # noqa: ARG001 return sentinel monkeypatch.setattr(emb_mod, "render_datasource_pair", stub) - captured: List[str] = [] + captured: list[str] = [] async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch signature - texts: List[str], *, model: Optional[str] = None, # noqa: ARG001 - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, # noqa: ARG001 + ) -> list[list[float] | None]: captured.extend(texts) return [[0.0, 1.0]] * len(texts) diff --git a/tests/test_search_service.py b/tests/test_search_service.py index 1c3fd2fe..994665cc 100644 --- a/tests/test_search_service.py +++ b/tests/test_search_service.py @@ -23,7 +23,7 @@ import asyncio import sqlite3 import tempfile -from typing import AsyncIterator, Tuple +from collections.abc import AsyncIterator import pytest import pytest_asyncio @@ -506,7 +506,7 @@ async def test_memory_hit_query_field_is_on_searchhit() -> None: @pytest_asyncio.fixture -async def stale_setup() -> AsyncIterator[Tuple[StorageBackend, SlayerQueryEngine]]: +async def stale_setup() -> AsyncIterator[tuple[StorageBackend, SlayerQueryEngine]]: """SQLite-backed storage + engine with a populated ``orders`` table AND a stored ``orders`` model whose ``status`` column has stale sample-value data (``sampled_values=None``).""" @@ -553,7 +553,7 @@ async def stale_setup() -> AsyncIterator[Tuple[StorageBackend, SlayerQueryEngine @pytest.mark.asyncio async def test_search_service_accepts_engine_kwarg( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], + stale_setup: tuple[StorageBackend, SlayerQueryEngine], ) -> None: """DEV-1516: SearchService now accepts an optional ``engine=`` kwarg so the search path can run the sample-refresh helper. Codex finding #1. @@ -582,7 +582,7 @@ async def test_search_service_engine_default_none_works( @pytest.mark.asyncio async def test_search_refreshes_stale_categorical_column_hit( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], + stale_setup: tuple[StorageBackend, SlayerQueryEngine], ) -> None: """DEV-1516 core contract: when search returns a stale categorical column hit, the helper fires, persists, and the hit's ``text`` @@ -614,7 +614,7 @@ async def test_search_refreshes_stale_categorical_column_hit( @pytest.mark.asyncio async def test_search_refreshes_stale_column_hit_via_question_corpus( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], + stale_setup: tuple[StorageBackend, SlayerQueryEngine], ) -> None: """Codex round-3 finding #2: refresh must also apply to column hits that surface via the question/corpus path (channels 2 & 3), not only the @@ -654,7 +654,7 @@ async def test_search_refreshes_stale_column_hit_via_question_corpus( @pytest.mark.asyncio async def test_search_no_refresh_when_engine_is_none( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], + stale_setup: tuple[StorageBackend, SlayerQueryEngine], ) -> None: """Without an engine the refresh hook is a no-op. The hit still surfaces (with stale text); storage stays untouched.""" @@ -679,7 +679,7 @@ async def test_search_no_refresh_when_engine_is_none( @pytest.mark.asyncio async def test_search_stale_text_preserved_when_profile_raises( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], monkeypatch, + stale_setup: tuple[StorageBackend, SlayerQueryEngine], monkeypatch, ) -> None: """If the helper's profile_column raises, the search hit falls back to the original (stale) rendered text. No crash. Storage stays untouched. @@ -738,7 +738,7 @@ async def tracking_persist(**kwargs): @pytest.mark.asyncio async def test_search_numeric_column_hit_not_refreshed( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], monkeypatch, + stale_setup: tuple[StorageBackend, SlayerQueryEngine], monkeypatch, ) -> None: """The helper is categorical-only. A numeric column hit must NOT trigger profile_column from the search refresh path. @@ -784,7 +784,7 @@ async def counting_profile(*, model, column, engine): @pytest.mark.asyncio async def test_search_per_model_serialization_concurrent_hits( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], monkeypatch, + stale_setup: tuple[StorageBackend, SlayerQueryEngine], monkeypatch, ) -> None: """Codex finding #2: storage.update_column_sampled does a model-level read-modify-write. Two stale-column hits on the SAME model must NOT @@ -883,7 +883,7 @@ async def tracking_persist(**kwargs): @pytest.mark.asyncio async def test_search_cross_model_concurrency_is_allowed( - stale_setup: Tuple[StorageBackend, SlayerQueryEngine], monkeypatch, + stale_setup: tuple[StorageBackend, SlayerQueryEngine], monkeypatch, ) -> None: """Codex round-3 finding #4: per the plan, refresh serializes within a ``(data_source, model_name)`` group AND parallelizes across groups via diff --git a/tests/test_search_service_dispatch.py b/tests/test_search_service_dispatch.py index 13c1478e..216898da 100644 --- a/tests/test_search_service_dispatch.py +++ b/tests/test_search_service_dispatch.py @@ -20,7 +20,7 @@ from __future__ import annotations import tempfile -from typing import AsyncIterator, List, Optional +from collections.abc import AsyncIterator import pytest_asyncio @@ -53,9 +53,9 @@ class _RecordingRetriever(Retriever): def __init__( self, *, name: str, - upsert_warn: Optional[str] = None, - subtree_warn: Optional[str] = None, - ds_warn: Optional[str] = None, + upsert_warn: str | None = None, + subtree_warn: str | None = None, + ds_warn: str | None = None, raise_on_upsert: bool = False, raise_on_subtree: bool = False, raise_on_datasource: bool = False, @@ -67,20 +67,20 @@ def __init__( self._raise_on_upsert = raise_on_upsert self._raise_on_subtree = raise_on_subtree self._raise_on_datasource = raise_on_datasource - self.upsert_calls: List[Memory] = [] - self.subtree_calls: List[SlayerModel] = [] - self.ds_calls: List[str] = [] + self.upsert_calls: list[Memory] = [] + self.subtree_calls: list[SlayerModel] = [] + self.ds_calls: list[str] = [] async def retrieve(self, **kwargs) -> RetrievalResult: return RetrievalResult() - async def upsert_memory(self, memory: Memory) -> List[str]: + async def upsert_memory(self, memory: Memory) -> list[str]: self.upsert_calls.append(memory) if self._raise_on_upsert: raise RuntimeError(f"{self.name} boom") return [self._upsert_warn] if self._upsert_warn else [] - async def refresh_model_subtree(self, model: SlayerModel) -> List[str]: + async def refresh_model_subtree(self, model: SlayerModel) -> list[str]: self.subtree_calls.append(model) if self._raise_on_subtree: raise RuntimeError(f"{self.name} subtree-boom") @@ -90,9 +90,9 @@ async def refresh_datasource( self, *, name: str, - models: List[SlayerModel], - description: Optional[str] = None, - ) -> List[str]: + models: list[SlayerModel], + description: str | None = None, + ) -> list[str]: self.ds_calls.append(name) if self._raise_on_datasource: raise RuntimeError(f"{self.name} ds-boom") diff --git a/tests/test_search_service_orchestration.py b/tests/test_search_service_orchestration.py index ad6f6b7a..a90ff289 100644 --- a/tests/test_search_service_orchestration.py +++ b/tests/test_search_service_orchestration.py @@ -28,7 +28,7 @@ import asyncio import tempfile -from typing import AsyncIterator, List, Optional, Set +from collections.abc import AsyncIterator from unittest.mock import patch import pytest_asyncio @@ -73,10 +73,10 @@ class _CapturingRetriever(Retriever): def __init__( self, *, name: str, - memory_ranking: Optional[List[str]] = None, - entity_ranking: Optional[List[str]] = None, - text_by_id: Optional[dict] = None, - warning: Optional[str] = None, + memory_ranking: list[str] | None = None, + entity_ranking: list[str] | None = None, + text_by_id: dict | None = None, + warning: str | None = None, delay_s: float = 0.0, ) -> None: self.name = name @@ -87,22 +87,22 @@ def __init__( self._delay_s = delay_s # Capture by identity (id()) so we can prove the SAME object # was passed to every retriever, not merely equal copies. - self.captured_valid_canonicals_ids: List[int] = [] - self.captured_valid_canonicals: List[Set[str]] = [] - self.captured_corpus_ids: List[int] = [] - self.captured_all_memories_ids: List[int] = [] - self.captured_all_memories_lens: List[int] = [] - self.captured_datasources: List[Optional[str]] = [] + self.captured_valid_canonicals_ids: list[int] = [] + self.captured_valid_canonicals: list[set[str]] = [] + self.captured_corpus_ids: list[int] = [] + self.captured_all_memories_ids: list[int] = [] + self.captured_all_memories_lens: list[int] = [] + self.captured_datasources: list[str | None] = [] self.retrieve_call_count = 0 async def retrieve( self, *, - query_entities: List[str], - question: Optional[str], - all_memories: List[Memory], + query_entities: list[str], + question: str | None, + all_memories: list[Memory], valid_canonicals: set, - corpus: Optional[Corpus], - datasource: Optional[str], + corpus: Corpus | None, + datasource: str | None, ) -> RetrievalResult: self.retrieve_call_count += 1 self.captured_valid_canonicals_ids.append(id(valid_canonicals)) diff --git a/tests/test_search_surfaces.py b/tests/test_search_surfaces.py index 5d6c077d..89f05307 100644 --- a/tests/test_search_surfaces.py +++ b/tests/test_search_surfaces.py @@ -11,7 +11,7 @@ import json import sys import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_search_three_channel.py b/tests/test_search_three_channel.py index 2f77faa5..79fda6ca 100644 --- a/tests/test_search_three_channel.py +++ b/tests/test_search_three_channel.py @@ -10,7 +10,7 @@ from __future__ import annotations import tempfile -from typing import Iterator, List, Optional +from collections.abc import Iterator import pytest @@ -114,9 +114,9 @@ async def test_question_with_embeddings_returns_entity( text_to_vec: dict = {} async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch async signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: - out: List[Optional[List[float]]] = [] + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: + out: list[list[float] | None] = [] for idx, t in enumerate(texts): # Distinguish "amount" column doc by a marker vector — every # other doc gets a flatter base vector. @@ -147,8 +147,8 @@ async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch async # Stub the query-side embedding: align with the "amount" column # vector so cosine ranks it #1. async def stub_embed_query( # NOSONAR(S7503) — stub matches embed_query async signature - text: str, *, model: Optional[str] = None, - ) -> List[float]: + text: str, *, model: str | None = None, + ) -> list[float]: # Align with amount column. return [1.0, 0.0, 0.0, 0.0] @@ -179,11 +179,11 @@ async def test_entity_hits_now_carry_rrf_fused_score( await _seed_basic_corpus(storage) async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch async signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: return [[0.0, 0.0, 0.0, 0.0] for _ in texts] - async def stub_embed_query(*_a, **_kw) -> List[float]: # NOSONAR(S7503) — stub matches embed_query async signature + async def stub_embed_query(*_a, **_kw) -> list[float]: # NOSONAR(S7503) — stub matches embed_query async signature return [0.0, 0.0, 0.0, 0.0] monkeypatch.setattr( @@ -225,8 +225,8 @@ async def test_query_embed_failure_warns_and_continues( await _seed_basic_corpus(storage) async def stub_embed_batch( # NOSONAR(S7503) — stub matches embed_batch async signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: return [[0.1, 0.1, 0.1, 0.1] for _ in texts] monkeypatch.setattr( @@ -267,9 +267,9 @@ async def test_entity_only_does_not_trigger_channel_3( learning="learning", entities=["dsx.orders"], query=None, ) - embed_called: List[str] = [] + embed_called: list[str] = [] - async def stub_embed_query(text: str, *, model: Optional[str] = None): # NOSONAR(S7503) — stub matches embed_query async signature + async def stub_embed_query(text: str, *, model: str | None = None): # NOSONAR(S7503) — stub matches embed_query async signature embed_called.append(text) return [0.1, 0.1, 0.1, 0.1] diff --git a/tests/test_search_unified.py b/tests/test_search_unified.py index 93023822..73253727 100644 --- a/tests/test_search_unified.py +++ b/tests/test_search_unified.py @@ -16,7 +16,7 @@ import json import sys import tempfile -from typing import AsyncIterator +from collections.abc import AsyncIterator import pytest import pytest_asyncio diff --git a/tests/test_sidecar_embedding_store.py b/tests/test_sidecar_embedding_store.py index a0ced312..a00cbad2 100644 --- a/tests/test_sidecar_embedding_store.py +++ b/tests/test_sidecar_embedding_store.py @@ -13,7 +13,7 @@ import sqlite3 import tempfile from datetime import datetime, timezone -from typing import Iterator, List +from collections.abc import Iterator import pytest @@ -38,7 +38,7 @@ def _embed( model: str = "openai/test-embedding", kind: EntityKind = "memory", text_hash: str = "h0", - vector: List[float] | None = None, + vector: list[float] | None = None, ) -> Embedding: return Embedding( canonical_id=canonical_id, @@ -166,7 +166,7 @@ async def test_save_many_empty_is_noop( from slayer.storage import sidecar_embedding_store as _mod real_connect = _mod.sqlite3.connect - connect_calls: List[tuple] = [] + connect_calls: list[tuple] = [] def _spy(*args, **kwargs): connect_calls.append(args) @@ -266,7 +266,7 @@ async def test_get_many_empty_input_returns_empty( from slayer.storage import sidecar_embedding_store as _mod real_connect = _mod.sqlite3.connect - connect_calls: List[tuple] = [] + connect_calls: list[tuple] = [] def _spy(*args, **kwargs): connect_calls.append(args) diff --git a/tests/test_sql_client_in_memory_async.py b/tests/test_sql_client_in_memory_async.py index 8a455027..d65ea799 100644 --- a/tests/test_sql_client_in_memory_async.py +++ b/tests/test_sql_client_in_memory_async.py @@ -20,7 +20,8 @@ """ import asyncio -from typing import Any, Iterable +from typing import Any +from collections.abc import Iterable import pytest import sqlalchemy as sa diff --git a/tests/test_startup_ingest.py b/tests/test_startup_ingest.py index a52e1e50..09ce8045 100644 --- a/tests/test_startup_ingest.py +++ b/tests/test_startup_ingest.py @@ -22,7 +22,6 @@ import tempfile import types from pathlib import Path -from typing import List, Optional from unittest.mock import AsyncMock, MagicMock import pytest @@ -59,9 +58,9 @@ def _ds(name: str) -> DatasourceConfig: def _stub_storage( *, - names: List[str], - list_raises: Optional[BaseException] = None, - missing: Optional[List[str]] = None, + names: list[str], + list_raises: BaseException | None = None, + missing: list[str] | None = None, ) -> MagicMock: """Build a storage stub with async `list_datasources` and `get_datasource`. @@ -74,7 +73,7 @@ def _stub_storage( else: storage.list_datasources = AsyncMock(return_value=list(names)) - async def _get(name: str) -> Optional[DatasourceConfig]: # NOSONAR(S7503) — must be `async def` to be a valid AsyncMock.side_effect + async def _get(name: str) -> DatasourceConfig | None: # NOSONAR(S7503) — must be `async def` to be a valid AsyncMock.side_effect if name in missing_set: return None if name in names: @@ -139,7 +138,7 @@ def _patch_ingester(monkeypatch, behaviour): `behaviour` is a callable `(name, ds) -> IdempotentIngestResult` (or `(name, ds) -> raise`). Awaited by the orchestrator. """ - call_log: List[str] = [] + call_log: list[str] = [] async def fake(*, datasource, storage, schema, include_tables, exclude_tables): # NOSONAR(S7503) — must be `async def` to replace the real async ingest_datasource_idempotent call_log.append(datasource.name) @@ -859,8 +858,8 @@ async def test_orchestrator_refreshes_memory_embeddings_for_each_datasource( ) async def fake_embed_batch( # NOSONAR(S7503) — must be `async def` to match the patched embed_batch signature - texts: List[str], *, model: Optional[str] = None, - ) -> List[Optional[List[float]]]: + texts: list[str], *, model: str | None = None, + ) -> list[list[float] | None]: return [[0.1, 0.2, 0.3] for _ in texts] monkeypatch.setattr( diff --git a/tests/test_tantivy_retriever.py b/tests/test_tantivy_retriever.py index 862c344f..17685017 100644 --- a/tests/test_tantivy_retriever.py +++ b/tests/test_tantivy_retriever.py @@ -14,7 +14,6 @@ from __future__ import annotations import time -from typing import List from unittest.mock import patch from slayer.core.enums import DataType @@ -42,7 +41,7 @@ def _model(name: str = "orders") -> SlayerModel: ) -def _memories() -> List[Memory]: +def _memories() -> list[Memory]: return [ Memory(id="1", learning="amount column is in cents", entities=["mydb.orders.amount"]), @@ -135,7 +134,7 @@ async def test_two_kind_filtered_queries_run_sequentially() -> None: """ retriever = TantivyRetriever() corpus = _build_corpus() - call_kind_filters: List[object] = [] + call_kind_filters: list[object] = [] def fake_search_index(*, index, question, limit, kind_filter=None, exclude_kind=None, fields=None): diff --git a/tests/test_validate_models.py b/tests/test_validate_models.py index 587647c2..0dea8b7b 100644 --- a/tests/test_validate_models.py +++ b/tests/test_validate_models.py @@ -19,7 +19,6 @@ import sqlite3 import tempfile from pathlib import Path -from typing import Dict, List, Optional import pytest @@ -56,10 +55,10 @@ def _orders_model( *, data_source: str = "ds", - columns: Optional[List[Column]] = None, - joins: Optional[List[ModelJoin]] = None, - measures: Optional[List[ModelMeasure]] = None, - filters: Optional[List[str]] = None, + columns: list[Column] | None = None, + joins: list[ModelJoin] | None = None, + measures: list[ModelMeasure] | None = None, + filters: list[str] | None = None, sql_table: str = "orders", ) -> SlayerModel: return SlayerModel( @@ -81,8 +80,8 @@ def _orders_model( def _live_orders( *, - columns: Optional[Dict[str, DataType]] = None, - pk_columns: Optional[set[str]] = None, + columns: dict[str, DataType] | None = None, + pk_columns: set[str] | None = None, ) -> LiveTable: return LiveTable( columns=columns @@ -96,7 +95,7 @@ def _live_orders( ) -def _entry_for(model_name: str, entries: List) -> Optional[object]: +def _entry_for(model_name: str, entries: list) -> object | None: for e in entries: if e.model_name == model_name: return e @@ -997,8 +996,8 @@ async def _setup_sensordata( workspace: Path, *, values: list, - columns: List[Column], - measures: Optional[List[ModelMeasure]] = None, + columns: list[Column], + measures: list[ModelMeasure] | None = None, ds_type: str = "sqlite", ) -> tuple[SlayerQueryEngine, str]: """Persist a sensordata model whose ``tempstabidx`` column is From da29827d493958debc6445cd21324277a3c5f4fa Mon Sep 17 00:00:00 2001 From: AivanF Date: Thu, 25 Jun 2026 14:06:13 +0300 Subject: [PATCH 2/2] Fixes for new content --- slayer/api/server.py | 6 +- slayer/client/slayer_client.py | 16 ++-- slayer/inspect/model_render.py | 126 +++++++++++++------------- slayer/inspect/service.py | 53 ++++++----- slayer/mcp/server.py | 160 ++++++++++++++++----------------- 5 files changed, 180 insertions(+), 181 deletions(-) diff --git a/slayer/api/server.py b/slayer/api/server.py index 381a812c..045687f1 100644 --- a/slayer/api/server.py +++ b/slayer/api/server.py @@ -181,8 +181,8 @@ class InspectRequest(BaseModel): format: str = "markdown" num_rows: int = 3 show_sql: bool = False - sections: Optional[List[str]] = None - descriptions_max_chars: Optional[int] = None + sections: list[str] | None = None + descriptions_max_chars: int | None = None def _slayer_version() -> str: @@ -728,7 +728,7 @@ async def search(request: SearchRequest) -> dict[str, Any]: } }, ) - async def inspect(request: InspectRequest) -> Dict[str, Any]: + async def inspect(request: InspectRequest) -> dict[str, Any]: try: result = await inspect_service.inspect( reference=request.reference, diff --git a/slayer/client/slayer_client.py b/slayer/client/slayer_client.py index 0463e1f8..45231f71 100644 --- a/slayer/client/slayer_client.py +++ b/slayer/client/slayer_client.py @@ -471,8 +471,8 @@ async def inspect( format: str = "markdown", num_rows: int = 3, show_sql: bool = False, - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, ) -> str: """Inspect EXACTLY one entity by reference and kind (DEV-1588). @@ -522,10 +522,10 @@ def _build_inspect_body( format: str, num_rows: int, show_sql: bool, - sections: Optional[List[str]], - descriptions_max_chars: Optional[int], - ) -> Dict[str, Any]: - body: Dict[str, Any] = { + sections: list[str] | None, + descriptions_max_chars: int | None, + ) -> dict[str, Any]: + body: dict[str, Any] = { "reference": reference, "entity_type": entity_type, "compact": compact, @@ -587,8 +587,8 @@ def inspect_sync( format: str = "markdown", num_rows: int = 3, show_sql: bool = False, - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, ) -> str: """Synchronous variant of :meth:`inspect` (DEV-1588).""" if self._storage is not None: diff --git a/slayer/inspect/model_render.py b/slayer/inspect/model_render.py index a88e7d93..d9ca70d0 100644 --- a/slayer/inspect/model_render.py +++ b/slayer/inspect/model_render.py @@ -14,7 +14,7 @@ import json import logging -from typing import Any, Dict, List, Optional, Tuple +from typing import Any import sqlalchemy as sa @@ -100,7 +100,7 @@ def _cell_is_present(value: Any) -> bool: return True -def _truncate_description(text: Optional[str], max_chars: Optional[int]) -> Optional[str]: +def _truncate_description(text: str | None, max_chars: int | None) -> str | None: """Trim a description to ``max_chars`` and append the truncation marker. Returns the input unchanged when ``max_chars`` is ``None`` or the text is @@ -114,7 +114,7 @@ def _truncate_description(text: Optional[str], max_chars: Optional[int]) -> Opti return text[:max_chars] + _TRUNCATION_MARKER -def _format_meta(meta: Optional[Dict[str, Any]]) -> Optional[str]: +def _format_meta(meta: dict[str, Any] | None) -> str | None: """Compact JSON for the ``inspect_model`` meta cell. Returns ``None`` when ``meta`` is ``None`` so ``_markdown_table``'s @@ -126,8 +126,8 @@ def _format_meta(meta: Optional[Dict[str, Any]]) -> Optional[str]: def _resolve_inspect_sections( - sections: Optional[List[str]], -) -> Tuple[List[str], List[str]]: + sections: list[str] | None, +) -> tuple[list[str], list[str]]: """Validate and normalise the ``sections`` argument for ``inspect_model``. Returns ``(resolved, unknown)`` where ``resolved`` is the list of valid @@ -155,11 +155,11 @@ def _resolve_inspect_sections( def _render_inspect_footer( *, - included: List[str], - names_only: List[str], - omitted: List[str], - unknown: List[str], -) -> Optional[str]: + included: list[str], + names_only: list[str], + omitted: list[str], + unknown: list[str], +) -> str | None: """Build the per-call truncation footer for ``inspect_model``. Returns ``None`` when there is nothing to report (no trimming, no @@ -167,7 +167,7 @@ def _render_inspect_footer( """ if not (names_only or omitted or unknown): return None - lines: List[str] = [] + lines: list[str] = [] if unknown: # repr() escapes newlines / quote chars so a caller-supplied value # like "foo\n> evil" can't forge additional footer lines. @@ -186,7 +186,7 @@ def _render_inspect_footer( return "\n".join(lines) if lines else None -def _markdown_table(rows: List[Dict[str, Any]], columns: List[str]) -> str: +def _markdown_table(rows: list[dict[str, Any]], columns: list[str]) -> str: """Render a list of row dicts as a GitHub-flavored markdown table. Columns with no present cell across every row are dropped automatically so @@ -227,11 +227,11 @@ def _markdown_table(rows: List[Dict[str, Any]], columns: List[str]) -> str: def _choose_sample_dims( model: SlayerModel, -) -> Tuple[List[Dict[str, str]], set]: +) -> tuple[list[dict[str, str]], set]: """Pick up to two categorical (TEXT/BOOLEAN) non-hidden, non-PK columns to group the sample by, so they aren't also aggregated as measures (count_distinct(status) grouped by status is always 1).""" - dims: List[Dict[str, str]] = [] + dims: list[dict[str, str]] = [] dim_names: set = set() for c in model.columns: if c.hidden or c.primary_key: @@ -249,8 +249,8 @@ def _choose_sample_dims( def _choose_sample_agg( column: Column, *, - measure_types: Dict[str, str], -) -> Optional[str]: + measure_types: dict[str, str], +) -> str | None: """Pick a sample aggregation for ``column``, or ``None`` to skip it. - With a restricted ``allowed_aggregations`` that excludes ``avg``: prefer @@ -279,8 +279,8 @@ def _choose_sample_agg( def _build_sample_query_args( model: SlayerModel, num_rows: int, - measure_types: Optional[Dict[str, str]] = None, -) -> Dict[str, Any]: + measure_types: dict[str, str] | None = None, +) -> dict[str, Any]: """Build the ``SlayerQuery`` payload for ``inspect_model``'s sample data. First measure is always ``*:count``; then one aggregation per non-hidden, @@ -289,7 +289,7 @@ def _build_sample_query_args( measure_types = measure_types or {} dims, dim_names = _choose_sample_dims(model) - measures: List[Dict[str, str]] = [{"formula": "*:count"}] + measures: list[dict[str, str]] = [{"formula": "*:count"}] for c in model.columns: if c.hidden or c.primary_key or c.name in dim_names: continue @@ -307,10 +307,10 @@ def _build_sample_query_args( def _strip_model_prefix( - columns: List[str], - data: List[Dict[str, Any]], + columns: list[str], + data: list[dict[str, Any]], model_name: str, -) -> Tuple[List[str], List[Dict[str, Any]]]: +) -> tuple[list[str], list[dict[str, Any]]]: """Drop the redundant ``{model_name}.`` prefix from sample-data column keys. Keeps the markdown table compact (the model name already appears in the @@ -328,7 +328,7 @@ def _strip(key: str) -> str: async def _get_row_count( model: SlayerModel, engine: SlayerQueryEngine, -) -> Optional[int]: +) -> int | None: """Return the total row count of ``model``'s underlying table, or ``None`` on any failure. Uses a bare ``*:count`` query — the same aggregation a user would run to ask for the count. @@ -360,7 +360,7 @@ async def _get_row_count( async def _collect_measure_profile( model: SlayerModel, engine: SlayerQueryEngine, -) -> Dict[str, str]: +) -> dict[str, str]: """Probe min/max for each non-hidden, non-primary-key NUMERIC/TEMPORAL column via a single batched query. @@ -394,7 +394,7 @@ async def _collect_measure_profile( "type": str(c.type)} for c in columns ] - measures_payload: List[Dict[str, str]] = [] + measures_payload: list[dict[str, str]] = [] for c in columns: measures_payload.append({"formula": f"_slayer_probe_{c.name}:min"}) measures_payload.append({"formula": f"_slayer_probe_{c.name}:max"}) @@ -409,7 +409,7 @@ async def _collect_measure_profile( except Exception: return {} - result: Dict[str, str] = {} + result: dict[str, str] = {} for c in columns: mn = row.get(f"{model.name}._slayer_probe_{c.name}_min") mx = row.get(f"{model.name}._slayer_probe_{c.name}_max") @@ -420,7 +420,7 @@ async def _collect_measure_profile( return result -def _build_backing_query_info(model: SlayerModel) -> Optional[dict]: +def _build_backing_query_info(model: SlayerModel) -> dict | None: """Build the ``backing_query`` block for inspect_model output. Returns ``None`` for non-query-backed models. For query-backed models, @@ -435,7 +435,7 @@ def _build_backing_query_info(model: SlayerModel) -> Optional[dict]: from slayer.core.query import extract_placeholder_names all_placeholders: set = set() - stage_dicts: List[dict] = [] + stage_dicts: list[dict] = [] # A placeholder is "required" only if it has no default at any layer the # engine consults: model.query_variables OR the stage's own variables. defaulted: set = set(model.query_variables.keys()) @@ -482,7 +482,7 @@ def _render_stage_field_list(key: str, val: list) -> str: return "; ".join(_render_field_value(v) for v in val) -def _render_source_model(src: Any) -> Optional[str]: +def _render_source_model(src: Any) -> str | None: """Render a stage's ``source_model`` (str or ModelExtension dict).""" if isinstance(src, str): return f"- source_model: `{src}`" @@ -493,10 +493,10 @@ def _render_source_model(src: Any) -> Optional[str]: return None -def _render_stage(i: int, stage: dict, total: int) -> List[str]: +def _render_stage(i: int, stage: dict, total: int) -> list[str]: """Render one stage's markdown lines.""" title = stage.get("name") or ("final" if i == total else f"stage {i}") - out: List[str] = [f"\n**{i}. {title}**"] + out: list[str] = [f"\n**{i}. {title}**"] src_line = _render_source_model(stage.get("source_model")) if src_line: out.append(src_line) @@ -510,7 +510,7 @@ def _render_stage(i: int, stage: dict, total: int) -> List[str]: def _backing_query_markdown_section(info: dict) -> str: """Format the ``backing_query`` info as a markdown section.""" - lines: List[str] = ["## Backing Query"] + lines: list[str] = ["## Backing Query"] stages = info.get("stages") or [] for i, stage in enumerate(stages, start=1): lines.extend(_render_stage(i, stage, len(stages))) @@ -541,8 +541,8 @@ def _source_type_for(model: SlayerModel) -> str: # --------------------------------------------------------------------------- def model_skeleton_fields( - *, model: SlayerModel, max_chars: Optional[int] = None, -) -> Dict[str, Any]: + *, model: SlayerModel, max_chars: int | None = None, +) -> dict[str, Any]: """Cheap, DB-free structured skeleton of a model. Shape: ``{name, canonical_id, description, column_names, measure_names, @@ -566,12 +566,12 @@ def model_skeleton_fields( } -def _skeleton_csv(names: List[str]) -> str: +def _skeleton_csv(names: list[str]) -> str: return ", ".join(names) if names else _NONE_PLACEHOLDER def render_model_skeleton( - *, model: SlayerModel, max_chars: Optional[int] = None, + *, model: SlayerModel, max_chars: int | None = None, ) -> str: """Heading-less markdown schema skeleton (DB-free). @@ -581,7 +581,7 @@ def render_model_skeleton( ``models_summary(compact)``). The caller prepends the ``#``/``##`` heading. """ fields = model_skeleton_fields(model=model, max_chars=max_chars) - lines: List[str] = [] + lines: list[str] = [] if fields["description"]: lines.append(fields["description"]) lines.append(f"Columns: {_skeleton_csv(fields['column_names'])}") @@ -595,12 +595,12 @@ async def render_model_inspection( # NOSONAR(S3776) — faithful extraction of *, model: SlayerModel, storage: StorageBackend, - engine: Optional[SlayerQueryEngine], + engine: SlayerQueryEngine | None, num_rows: int = 3, show_sql: bool = False, format: str = "markdown", - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, compact: bool = True, ) -> str: """Render a complete-yet-compact view of an already-resolved model. @@ -638,12 +638,12 @@ async def render_model_inspection( # NOSONAR(S3776) — faithful extraction of ] truncated_model_desc = _truncate_description(model.description, descriptions_max_chars) - out_sections: List[str] = [f"# Model: `{model.name}`"] + out_sections: list[str] = [f"# Model: `{model.name}`"] if truncated_model_desc: out_sections.append(truncated_model_desc) # Metadata bullets (incl. row_count from a cheap *:count query) - meta: List[str] = [] + meta: list[str] = [] if model.data_source: meta.append(f"- **data_source:** `{model.data_source}`") if model.sql_table: @@ -656,7 +656,7 @@ async def render_model_inspection( # NOSONAR(S3776) — faithful extraction of meta.append("- **hidden:** true") if model.meta is not None: meta.append(f"- **meta:** {json.dumps(model.meta, sort_keys=True, default=str)}") - row_count: Optional[int] = None + row_count: int | None = None if engine is not None: row_count = await _get_row_count(model=model, engine=engine) if row_count is not None: @@ -687,12 +687,12 @@ async def render_model_inspection( # NOSONAR(S3776) — faithful extraction of # DB-hitting computations — skip when their consumers aren't requested # (and when no engine is available, DEV-1588). # ------------------------------------------------------------------ - profile_by_name: Dict[str, str] = {} - profile_values_by_name: Dict[str, Optional[List[str]]] = {} - distinct_count_by_name: Dict[str, Optional[int]] = {} - measure_profile: Dict[str, str] = {} + profile_by_name: dict[str, str] = {} + profile_values_by_name: dict[str, list[str] | None] = {} + distinct_count_by_name: dict[str, int | None] = {} + measure_profile: dict[str, str] = {} if engine is not None and "columns" in included_set: - uncached_columns: List[Column] = [] + uncached_columns: list[Column] = [] for c in model.columns: if c.hidden or c.primary_key: continue @@ -734,9 +734,9 @@ async def render_model_inspection( # NOSONAR(S3776) — faithful extraction of async def _persist_sample( *, col_name: str, - sampled: Optional[str], - sampled_values: Optional[List[str]], - distinct_count: Optional[int], + sampled: str | None, + sampled_values: list[str] | None, + distinct_count: int | None, ) -> None: try: await storage.update_column_sampled( @@ -834,7 +834,7 @@ async def _persist_sample( # ``measure_types`` informs the sample query's choice of avg vs # count_distinct. Only needed when ``samples`` is in the included set. - measure_types: Dict[str, str] = {} + measure_types: dict[str, str] = {} if engine is not None and "samples" in included_set: measure_types = await engine.get_column_types( model_name=model.name, @@ -846,7 +846,7 @@ async def _persist_sample( # ------------------------------------------------------------------ visible_columns = [c for c in model.columns if not c.hidden] if "columns" in included_set: - col_rows: List[Dict[str, Any]] = [] + col_rows: list[dict[str, Any]] = [] for c in visible_columns: aggs = ", ".join(c.allowed_aggregations) if c.allowed_aggregations else "all" # DEV-1480: key-presence check (not ``or`` truthiness) so an @@ -889,7 +889,7 @@ async def _persist_sample( # Measures section # ------------------------------------------------------------------ if "measures" in included_set: - measure_rows: List[Dict[str, Any]] = [] + measure_rows: list[dict[str, Any]] = [] for mm in model.measures: measure_rows.append({ "name": mm.name, @@ -916,7 +916,7 @@ async def _persist_sample( # ------------------------------------------------------------------ if "aggregations" in included_set: if model.aggregations: - agg_rows: List[Dict[str, Any]] = [] + agg_rows: list[dict[str, Any]] = [] for a in model.aggregations: if a.params: if show_sql: @@ -951,7 +951,7 @@ async def _persist_sample( # Joins section # ------------------------------------------------------------------ if "joins" in included_set: - join_rows: List[Dict[str, Any]] = [] + join_rows: list[dict[str, Any]] = [] for j in model.joins: pairs = "; ".join(f"{src} = {tgt}" for src, tgt in j.join_pairs) join_rows.append({ @@ -974,9 +974,9 @@ async def _persist_sample( # ------------------------------------------------------------------ # Sample data (fully omitted when not in sections / no engine) # ------------------------------------------------------------------ - sample_sql: Optional[str] = None - sample_data: Optional[Dict[str, Any]] = None - sample_error: Optional[str] = None + sample_sql: str | None = None + sample_data: dict[str, Any] | None = None + sample_error: str | None = None if engine is not None and "samples" in included_set: query_args = _build_sample_query_args( model=model, num_rows=num_rows, measure_types=measure_types, @@ -1021,8 +1021,8 @@ async def _persist_sample( # ``None``; query-bearing memories are recall-only. Auto-pruned when # no learning-shaped memory matches. # ------------------------------------------------------------------ - relevant_learnings: List[Any] = [] - wanted: List[str] = [] + relevant_learnings: list[Any] = [] + wanted: list[str] = [] if "learnings" in included_set: ds = model.data_source wanted = [f"{ds}.{model.name}"] @@ -1070,7 +1070,7 @@ async def _persist_sample( ) if fmt == "json": - payload: Dict[str, Any] = { + payload: dict[str, Any] = { "model_name": model.name, "description": truncated_model_desc, "data_source": model.data_source, @@ -1092,7 +1092,7 @@ async def _persist_sample( # Columns if "columns" in included_set: - col_payloads: List[Dict[str, Any]] = [] + col_payloads: list[dict[str, Any]] = [] for c in visible_columns: # DEV-1480 key-presence (not ``or`` truthiness) so empty # string ``sampled=""`` (all-NULL categorical) survives. diff --git a/slayer/inspect/service.py b/slayer/inspect/service.py index 2b3e9233..1260af13 100644 --- a/slayer/inspect/service.py +++ b/slayer/inspect/service.py @@ -14,7 +14,6 @@ from __future__ import annotations import json -from typing import List, Optional from slayer.core.errors import ( AmbiguousModelError, @@ -68,7 +67,7 @@ def __init__( self, *, storage: StorageBackend, - engine: Optional["SlayerQueryEngine"] = None, + engine: SlayerQueryEngine | None = None, ) -> None: self._storage = storage self._engine = engine @@ -86,8 +85,8 @@ async def inspect( # NOSONAR(S3776) — single linear dispatch over the six ent format: str = "markdown", num_rows: int = 3, show_sql: bool = False, - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, ) -> str: # 1. Argument validation (raise ValueError). if entity_type not in VALID_ENTITY_TYPES: @@ -107,7 +106,7 @@ async def inspect( # NOSONAR(S3776) — single linear dispatch over the six ent ) # 2. Model-only-arg warnings (skip entirely for model entity_type). - warnings: List[str] = self._model_only_arg_warnings( + warnings: list[str] = self._model_only_arg_warnings( entity_type=entity_type, num_rows=num_rows, show_sql=show_sql, @@ -151,11 +150,11 @@ def _model_only_arg_warnings( entity_type: str, num_rows: int, show_sql: bool, - sections: Optional[List[str]], - ) -> List[str]: + sections: list[str] | None, + ) -> list[str]: if entity_type == "model": return [] - out: List[str] = [] + out: list[str] = [] # num_rows: warns for all non-model kinds when != default. if num_rows != 3: out.append(_warn_line(arg="num_rows", entity_type=entity_type)) @@ -174,7 +173,7 @@ def _model_only_arg_warnings( @staticmethod def _truncate_description_field( - text: str, max_chars: Optional[int], + text: str, max_chars: int | None, ) -> str: """Truncate only the ``Description: `` line(s) of a rendered entity blob — NOT the whole render. Mirrors ``inspect_model``'s @@ -183,7 +182,7 @@ def _truncate_description_field( survive a small ``descriptions_max_chars``.""" if max_chars is None: return text - out: List[str] = [] + out: list[str] = [] for line in text.split("\n"): if line.startswith(_DESCRIPTION_PREFIX): value = line[len(_DESCRIPTION_PREFIX):] @@ -197,7 +196,7 @@ def _truncate_description_field( return "\n".join(out) @staticmethod - def _markdown_with_warnings(body: str, warnings: List[str]) -> str: + def _markdown_with_warnings(body: str, warnings: list[str]) -> str: if not warnings: return body warn_block = "\n".join(f"> Warning: {w}" for w in warnings) @@ -215,8 +214,8 @@ async def _inspect_memory( reference: str, compact: bool, fmt: str, - descriptions_max_chars: Optional[int], - warnings: List[str], + descriptions_max_chars: int | None, + warnings: list[str], ) -> str: if not reference.startswith("memory:"): return ( @@ -281,8 +280,8 @@ async def _inspect_datasource( reference: str, compact: bool, fmt: str, - descriptions_max_chars: Optional[int], - warnings: List[str], + descriptions_max_chars: int | None, + warnings: list[str], ) -> str: try: res = await resolve_entity( @@ -302,7 +301,7 @@ async def _inspect_datasource( canonical = res.canonical_forms[0] known = set(await self._storage.list_datasources()) - ds_name: Optional[str] = None + ds_name: str | None = None if "." not in canonical and canonical in known: ds_name = canonical elif reference in known: @@ -324,8 +323,8 @@ async def _render_datasource( ds_name: str, compact: bool, fmt: str, - descriptions_max_chars: Optional[int], - warnings: List[str], + descriptions_max_chars: int | None, + warnings: list[str], ) -> str: cfg = await self._storage.get_datasource(ds_name) description = cfg.description if cfg is not None else None @@ -368,7 +367,7 @@ async def _render_datasource( "warnings": warnings, }, indent=2, default=str) - md_lines: List[str] = [f"Datasource: {ds_name}"] + md_lines: list[str] = [f"Datasource: {ds_name}"] if trunc_desc: md_lines.append(f"Description: {trunc_desc}") for m in models: @@ -392,9 +391,9 @@ async def _inspect_model( fmt: str, num_rows: int, show_sql: bool, - sections: Optional[List[str]], - descriptions_max_chars: Optional[int], - warnings: List[str], + sections: list[str] | None, + descriptions_max_chars: int | None, + warnings: list[str], ) -> str: try: canonical = await self._resolve_model_canonical(reference) @@ -455,7 +454,7 @@ async def _inspect_model( return json.dumps(payload, indent=2, default=str) return self._markdown_with_warnings(rendered, warnings) - async def _resolve_model_canonical(self, reference: str) -> Optional[str]: + async def _resolve_model_canonical(self, reference: str) -> str | None: """Resolve ``reference`` to a 2-part ``.`` canonical id, applying the Case-D entity_type=model override (a resolver that picked a datasource for a name that is also a model).""" @@ -500,8 +499,8 @@ async def _inspect_leaf( entity_type: str, compact: bool, fmt: str, - descriptions_max_chars: Optional[int], - warnings: List[str], + descriptions_max_chars: int | None, + warnings: list[str], ) -> str: try: res = await resolve_entity( @@ -558,8 +557,8 @@ def _render_leaf_entry( entity_type: str, compact: bool, fmt: str, - descriptions_max_chars: Optional[int], - warnings: List[str], + descriptions_max_chars: int | None, + warnings: list[str], ) -> str: trunc_desc = _truncate_description( text=entry.description, max_chars=descriptions_max_chars, diff --git a/slayer/mcp/server.py b/slayer/mcp/server.py index 9218ad2c..7811f1d7 100644 --- a/slayer/mcp/server.py +++ b/slayer/mcp/server.py @@ -2,7 +2,7 @@ import json import logging -from typing import Any, Dict, List, Optional, Tuple +from typing import Any import sqlalchemy as sa @@ -90,8 +90,8 @@ def _get_schemas(ds: DatasourceConfig) -> list[str]: def _fetch_tables( - ds: DatasourceConfig, schema_name: Optional[str] = None, -) -> Tuple[Optional[List[str]], Optional[str]]: + ds: DatasourceConfig, schema_name: str | None = None, +) -> tuple[list[str] | None, str | None]: """Inspect a datasource's table names. Returns ``(tables, None)`` on success or ``(None, friendly_error_message)`` @@ -121,7 +121,7 @@ def _empty_ingest_message(*, schema_name: str, ds: DatasourceConfig) -> str: return "\n".join(lines) -def _render_new_models_section(new_models: List[Any]) -> List[str]: +def _render_new_models_section(new_models: list[Any]) -> list[str]: if not new_models: return [] lines = [f"Created {len(new_models)} new model(s):"] @@ -132,7 +132,7 @@ def _render_new_models_section(new_models: List[Any]) -> List[str]: return lines -def _render_updated_section(updated: List[Any]) -> List[str]: +def _render_updated_section(updated: list[Any]) -> list[str]: if not updated: return [] lines = [f"Updated {len(updated)} existing model(s):"] @@ -146,7 +146,7 @@ def _render_updated_section(updated: List[Any]) -> List[str]: return lines -def _render_unchanged_section(unchanged: List[Any]) -> List[str]: +def _render_unchanged_section(unchanged: list[Any]) -> list[str]: if not unchanged: return [] return [ @@ -155,7 +155,7 @@ def _render_unchanged_section(unchanged: List[Any]) -> List[str]: ] -def _render_drift_section(to_delete: List[Any]) -> List[str]: +def _render_drift_section(to_delete: list[Any]) -> list[str]: if not to_delete: return [] out = ["", "Pending drift (run validate_models / apply manually):"] @@ -163,7 +163,7 @@ def _render_drift_section(to_delete: List[Any]) -> List[str]: return out -def _render_errors_section(errors: List[Any]) -> List[str]: +def _render_errors_section(errors: list[Any]) -> list[str]: if not errors: return [] out = ["", f"Errors ({len(errors)}):"] @@ -199,7 +199,7 @@ def _render_ingest_result( if not a.created and not a.new_columns and not a.new_joins ] - lines: List[str] = [] + lines: list[str] = [] lines.extend(_render_new_models_section(new_models)) lines.extend(_render_updated_section(updated)) lines.extend(_render_unchanged_section(unchanged)) @@ -287,25 +287,25 @@ def create_mcp_server( # NOSONAR(S3776) — FastMCP tool-registration factory; ) @mcp.tool(description=_help_description) - async def help(topic: Optional[str] = None) -> str: # noqa: A001 — intentional shadow of builtin inside factory + async def help(topic: str | None = None) -> str: # noqa: A001 — intentional shadow of builtin inside factory return render_help(topic=topic) @mcp.tool() async def query( # NOSONAR S107 — FastMCP introspects this signature to expose each query option as a typed MCP tool argument; collapsing into a dict would degrade the agent-facing schema source_model: str | ModelExtension | SlayerModel, - measures: Optional[List[Dict[str, str]]] = None, - dimensions: Optional[List[str]] = None, - filters: Optional[List[str]] = None, - time_dimensions: Optional[List[Dict[str, Any]]] = None, - order: Optional[List[Dict[str, str]]] = None, - limit: Optional[int] = None, - offset: Optional[int] = None, + measures: list[dict[str, str]] | None = None, + dimensions: list[str] | None = None, + filters: list[str] | None = None, + time_dimensions: list[dict[str, Any]] | None = None, + order: list[dict[str, str]] | None = None, + limit: int | None = None, + offset: int | None = None, whole_periods_only: bool = False, show_sql: bool = False, dry_run: bool = False, explain: bool = False, format: str = "markdown", - variables: Optional[Dict[str, Any]] = None, + variables: dict[str, Any] | None = None, distinct_dimension_values: bool = True, ) -> str: """Query data from a semantic model. Call inspect(reference=".", entity_type="model") first to see available columns and measures. @@ -347,7 +347,7 @@ async def query( # NOSONAR S107 — FastMCP introspects this signature to expos Before calling this tool, run ``search`` first, supplying the entities you're thinking of using (and/or the query itself via the ``query`` arg, or a free-text ``question``). Read the returned memories and consider any matching example queries before formulating the final query. """ - data: Dict[str, Any] = {"source_model": source_model} + data: dict[str, Any] = {"source_model": source_model} if dimensions: data["dimensions"] = list(dimensions) if filters: @@ -438,8 +438,8 @@ async def query( # NOSONAR S107 — FastMCP introspects this signature to expos @mcp.tool() async def query_nested( - queries: List[Dict[str, Any]], - variables: Optional[Dict[str, Any]] = None, + queries: list[dict[str, Any]], + variables: dict[str, Any] | None = None, show_sql: bool = False, dry_run: bool = False, explain: bool = False, @@ -560,7 +560,7 @@ async def models_summary( return f"Datasource '{datasource_name}' not found." all_names = await storage.list_models(data_source=datasource_name) - matched: List[SlayerModel] = [] + matched: list[SlayerModel] = [] for n in all_names: try: m = await storage.get_model(n, data_source=datasource_name) @@ -621,11 +621,11 @@ async def models_summary( indent=2, ) - sections: List[str] = [ + sections: list[str] = [ f"# Datasource: `{datasource_name}` — {len(matched)} model(s)" ] for m in matched: - model_lines: List[str] = [f"## `{m.name}`"] + model_lines: list[str] = [f"## `{m.name}`"] if m.description: model_lines.append(m.description) @@ -684,9 +684,9 @@ async def inspect_model( num_rows: int = 3, show_sql: bool = False, format: str = "markdown", - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, - data_source: Optional[str] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, + data_source: str | None = None, compact: bool = True, ) -> str: """DEPRECATED: use the ``inspect`` tool. Return a complete-yet-compact view of a semantic model. @@ -772,8 +772,8 @@ async def inspect( format: str = "markdown", num_rows: int = 3, show_sql: bool = False, - sections: Optional[List[str]] = None, - descriptions_max_chars: Optional[int] = None, + sections: list[str] | None = None, + descriptions_max_chars: int | None = None, ) -> str: """Inspect EXACTLY one entity by reference and kind. @@ -829,14 +829,14 @@ async def inspect( @mcp.tool() async def create_model( name: str, - sql_table: Optional[str] = None, - sql: Optional[str] = None, - data_source: Optional[str] = None, - description: Optional[str] = None, - columns: Optional[List[Dict[str, Any]]] = None, - measures: Optional[List[Dict[str, Any]]] = None, - query: Optional[Any] = None, - variables: Optional[Dict[str, Any]] = None, + sql_table: str | None = None, + sql: str | None = None, + data_source: str | None = None, + description: str | None = None, + columns: list[dict[str, Any]] | None = None, + measures: list[dict[str, Any]] | None = None, + query: Any | None = None, + variables: dict[str, Any] | None = None, ) -> str: """Create a new semantic model, either from a database table or from a query. @@ -935,7 +935,7 @@ def _upsert_entity( id_field: str, changes: list, label: str, - ) -> Optional[str]: + ) -> str | None: """Upsert a named entity in *entity_list*. Returns an error string on validation failure, ``None`` on success. @@ -970,23 +970,23 @@ def _upsert_entity( @mcp.tool() async def edit_model( model_name: str, - description: Optional[str] = None, - data_source: Optional[str] = None, - new_data_source: Optional[str] = None, - default_time_dimension: Optional[str] = None, - sql_table: Optional[str] = None, - sql: Optional[str] = None, - source_queries: Optional[List[Dict[str, Any]]] = None, + description: str | None = None, + data_source: str | None = None, + new_data_source: str | None = None, + default_time_dimension: str | None = None, + sql_table: str | None = None, + sql: str | None = None, + source_queries: list[dict[str, Any]] | None = None, query_variables: Any = _UNSET, - hidden: Optional[bool] = None, - columns: Optional[List[Dict[str, Any]]] = None, - measures: Optional[List[Dict[str, Any]]] = None, - aggregations: Optional[List[Dict[str, Any]]] = None, - joins: Optional[List[Dict[str, Any]]] = None, - add_filters: Optional[List[str]] = None, - remove_filters: Optional[List[str]] = None, - remove: Optional[Dict[str, List[str]]] = None, - meta: Optional[Dict[str, Any]] = _UNSET, + hidden: bool | None = None, + columns: list[dict[str, Any]] | None = None, + measures: list[dict[str, Any]] | None = None, + aggregations: list[dict[str, Any]] | None = None, + joins: list[dict[str, Any]] | None = None, + add_filters: list[str] | None = None, + remove_filters: list[str] | None = None, + remove: dict[str, list[str]] | None = None, + meta: dict[str, Any] | None = _UNSET, ) -> str: """Edit an existing model in a single call — update metadata, upsert columns/measures/aggregations/joins, manage filters, and remove entities. @@ -1053,7 +1053,7 @@ async def edit_model( return f"Model '{model_name}' not found." original_data_source = model.data_source - changes: List[str] = [] + changes: list[str] = [] # DEV-1375: track refresh-triggering changes so the post-save hook # knows whether to refresh just the touched columns or every # column on the model. @@ -1309,7 +1309,7 @@ async def edit_model( # changed the indexed text. Best-effort: any raise here is # captured into ``refresh_warnings`` so the save's success # status survives a flaky embedding API. - refresh_warnings: List[str] = [] + refresh_warnings: list[str] = [] if changed_columns or model_level_change or model_doc_changed: try: refresh_warnings = await handle_edit_refresh( @@ -1346,13 +1346,13 @@ async def edit_model( async def create_datasource( name: str, type: str, - host: Optional[str] = None, - port: Optional[int] = None, - database: Optional[str] = None, - username: Optional[str] = None, - password: Optional[str] = None, - connection_string: Optional[str] = None, - schema_name: Optional[str] = None, + host: str | None = None, + port: int | None = None, + database: str | None = None, + username: str | None = None, + password: str | None = None, + connection_string: str | None = None, + schema_name: str | None = None, auto_ingest: bool = True, ) -> str: """Create a database connection, verify it, and auto-ingest models. Use ${ENV_VAR} syntax in credentials to reference environment variables. @@ -1512,7 +1512,7 @@ async def describe_datasource( @mcp.tool() async def edit_datasource( name: str, - description: Optional[str] = None, + description: str | None = None, ) -> str: """Update a datasource's metadata. @@ -1540,9 +1540,9 @@ async def edit_datasource( # round-7 review: the refresh is post-save and best-effort — # log a warning if it raises and surface a partial-success # message rather than telling the agent the save itself failed. - refresh_warning: Optional[str] = None + refresh_warning: str | None = None if description is not None and description != old_description: - models_in_ds: List[SlayerModel] = [] + models_in_ds: list[SlayerModel] = [] for model_name in await storage.list_models(data_source=name): m = await storage.get_model(model_name, data_source=name) if m is not None: @@ -1570,7 +1570,7 @@ async def edit_datasource( # ----------------------------------------------------------------------- @mcp.tool() - async def delete_model(name: str, data_source: Optional[str] = None) -> str: + async def delete_model(name: str, data_source: str | None = None) -> str: """Delete a semantic model. Args: @@ -1588,7 +1588,7 @@ async def delete_model(name: str, data_source: Optional[str] = None) -> str: return f"Model '{name}' not found." @mcp.tool() - async def validate_models(data_source: Optional[str] = None) -> str: + async def validate_models(data_source: str | None = None) -> str: """Diff persisted SLayer models against the live database schema(s). Returns a JSON-serialized list of pending delete operations @@ -1668,7 +1668,7 @@ async def ingest_datasource_models(datasource_name: str, include_tables: str = " ) @mcp.tool() - async def set_datasource_priority(priority: List[str]) -> str: + async def set_datasource_priority(priority: list[str]) -> str: """Configure how SLayer disambiguates bare model names that exist in multiple datasources. @@ -1714,8 +1714,8 @@ def _format_resolution_error(exc: Exception) -> str: async def save_memory( learning: str, linked_entities: Any, - id: Optional[str] = None, # noqa: A002 — MCP arg name - description: Optional[str] = None, + id: str | None = None, # noqa: A002 — MCP arg name + description: str | None = None, ) -> str: """Save an agent memory: a free-form note plus the SLayer entities it concerns. @@ -1824,12 +1824,12 @@ async def forget_memory(id: Any) -> str: # noqa: A002 — MCP arg name @mcp.tool() async def search( - entities: Optional[List[str]] = None, + entities: list[str] | None = None, query: Any = None, - question: Optional[str] = None, - datasource: Optional[str] = None, + question: str | None = None, + datasource: str | None = None, max_results: int = 10, - cypher_filter: Optional[str] = None, + cypher_filter: str | None = None, compact: bool = True, ) -> str: """Up to three-channel semantic search over memories + canonical entities. @@ -1902,12 +1902,12 @@ async def search( return mcp -def _build_dict(**kwargs: Any) -> Dict[str, Any]: +def _build_dict(**kwargs: Any) -> dict[str, Any]: """Build a dict from keyword arguments, excluding None values.""" return {k: v for k, v in kwargs.items() if v is not None} -def _format_table(data: List[Dict[str, Any]], columns: List[str], max_rows: int = 50) -> str: +def _format_table(data: list[dict[str, Any]], columns: list[str], max_rows: int = 50) -> str: """Format data as a pipe-separated table (used for sample data display).""" if not data: return "No results." @@ -1927,14 +1927,14 @@ def _format_table(data: List[Dict[str, Any]], columns: List[str], max_rows: int return result -def _format_json(data: List[Dict[str, Any]], columns: List[str]) -> str: +def _format_json(data: list[dict[str, Any]], columns: list[str]) -> str: """Format data as JSON array.""" import json return json.dumps(data, default=str) -def _format_csv(data: List[Dict[str, Any]], columns: List[str]) -> str: +def _format_csv(data: list[dict[str, Any]], columns: list[str]) -> str: """Format data as CSV.""" if not data: return "" @@ -1959,7 +1959,7 @@ def _format_output(result: SlayerResponse, fmt: str) -> str: return _format_json(data=result.data, columns=result.columns) -def _format_field_meta(entries: Dict[str, Any]) -> List[str]: +def _format_field_meta(entries: dict[str, Any]) -> list[str]: """Format a dict of field metadata entries into lines.""" lines = [] for col, fm in entries.items():