Skip to content

Commit

Permalink
Fix LogicalColumn mypy error
Browse files Browse the repository at this point in the history
Per Jim, Timespan columns are not expected at this point in the code and would have been throwing an error already if they did show up.
  • Loading branch information
dhirving committed Dec 28, 2023
1 parent 3c90ea0 commit 270740f
Showing 1 changed file with 17 additions and 5 deletions.
22 changes: 17 additions & 5 deletions python/lsst/daf/butler/registry/queries/butler_sql_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,14 +191,16 @@ def to_payload(self, relation: Relation) -> sql.Payload[LogicalColumn]:
# contains (at the same time, since they have the same columns,
# aside from the special 'rownum' window-function column).
search_columns = self.extract_mapping(target.columns, search.columns)
partition_by = [search_columns[tag] for tag in operation.dimensions]
partition_by = [
_assert_column_is_directly_usable_by_sqlalchemy(search_columns[tag])
for tag in operation.dimensions
]
row_number = sqlalchemy.sql.func.row_number()
rank_column = _assert_column_is_directly_usable_by_sqlalchemy(search_columns[operation.rank])
if partition_by:
rownum_column = row_number.over(
partition_by=partition_by, order_by=search_columns[operation.rank]
)
rownum_column = row_number.over(partition_by=partition_by, order_by=rank_column)
else:
rownum_column = row_number.over(order_by=search_columns[operation.rank])
rownum_column = row_number.over(order_by=rank_column)
window = self.select_items(
search_columns.items(), search, rownum_column.label("rownum")
).subquery(f"{operation.rank.dataset_type}_window")
Expand All @@ -209,3 +211,13 @@ def to_payload(self, relation: Relation) -> sql.Payload[LogicalColumn]:
)
case _:
return super().to_payload(relation)


def _assert_column_is_directly_usable_by_sqlalchemy(column: LogicalColumn) -> sqlalchemy.sql.ColumnElement:
"""Narrow a `LogicalColumn` to a SqlAlchemy ColumnElement, to satisfy the
typechecker in cases where no timespans are expected.
"""
if isinstance(column, TimespanDatabaseRepresentation):
raise TypeError("Timespans not expected here.")

Check warning on line 221 in python/lsst/daf/butler/registry/queries/butler_sql_engine.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/registry/queries/butler_sql_engine.py#L221

Added line #L221 was not covered by tests

return column

0 comments on commit 270740f

Please sign in to comment.