Skip to content

Commit

Permalink
Allow numpy integers to be used as literals
Browse files Browse the repository at this point in the history
Fix an issue where the new query system was rejecting numpy integers used in data IDs or bind values.
  • Loading branch information
dhirving committed Oct 16, 2024
1 parent 3abf212 commit 4ac578f
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 0 deletions.
4 changes: 4 additions & 0 deletions python/lsst/daf/butler/queries/tree/_column_literal.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
)

import datetime
import numbers
import uuid
import warnings
from base64 import b64decode, b64encode
Expand Down Expand Up @@ -400,6 +401,9 @@ def make_column_literal(value: LiteralValue) -> ColumnLiteral:
dec = icrs.dec.degree
lon_lat = lsst.sphgeom.LonLat.fromDegrees(ra, dec)
return _make_region_literal_from_lonlat(lon_lat)
case numbers.Integral():
# numpy.int64 and other integer-like values.
return IntColumnLiteral.from_value(int(value))

raise TypeError(f"Invalid type {type(value).__name__!r} of value {value!r} for column literal.")

Expand Down
25 changes: 25 additions & 0 deletions python/lsst/daf/butler/tests/butler_queries.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import astropy.coordinates
import astropy.time
from lsst.sphgeom import LonLat, Region
from numpy import int64

from .._butler import Butler
from .._collection_type import CollectionType
Expand Down Expand Up @@ -1964,6 +1965,30 @@ def test_default_data_id(self) -> None:
names = [x.name for x in result]
self.assertCountEqual(names, ["Cam2-G"])

def test_unusual_column_literals(self) -> None:
butler = self.make_butler("base.yaml")

# Users frequently use numpy integer types as literals in queries.
result = butler.query_dimension_records(
"detector", data_id={"instrument": "Cam1", "detector": int64(1)}
)
names = [x.full_name for x in result]
self.assertEqual(names, ["Aa"])

result = butler.query_dimension_records(
"detector", where="instrument='Cam1' and detector=an_integer", bind={"an_integer": int64(2)}
)
names = [x.full_name for x in result]
self.assertEqual(names, ["Ab"])

with butler.query() as query:
x = query.expression_factory
result = list(
query.dimension_records("detector").where(x.instrument == "Cam1", x.detector == int64(3))
)
names = [x.full_name for x in result]
self.assertEqual(names, ["Ba"])


def _get_exposure_ids_from_dimension_records(dimension_records: Iterable[DimensionRecord]) -> list[int]:
output = []
Expand Down

0 comments on commit 4ac578f

Please sign in to comment.