From 11799aa9d6848cfe7aec88ea7279e30dc78a9c5d Mon Sep 17 00:00:00 2001 From: gurcuff91 Date: Thu, 30 May 2024 10:49:32 -0400 Subject: [PATCH] Added geo queries --- mongotoy/db.py | 2 +- mongotoy/documents.py | 17 ++++---- mongotoy/expressions.py | 87 +++++++++++++++++++++++++++++++++++++++++ mongotoy/geodata.py | 21 ++++++---- pyproject.toml | 2 +- 5 files changed, 110 insertions(+), 19 deletions(-) diff --git a/mongotoy/db.py b/mongotoy/db.py index 0b8e833..bda5f9e 100644 --- a/mongotoy/db.py +++ b/mongotoy/db.py @@ -156,7 +156,7 @@ def _get_document_indexes( ( mappers.MultiPointMapper, mappers.MultiLineStringMapper, - mappers.PolygonMapper, + mappers.MultiPolygonMapper, ) ): indexes.append( diff --git a/mongotoy/documents.py b/mongotoy/documents.py index 560ead8..eee90e2 100644 --- a/mongotoy/documents.py +++ b/mongotoy/documents.py @@ -47,26 +47,23 @@ def __new__(mcls, name, bases, namespace, **kwargs): _fields.update(getattr(base, '__fields__', {})) # Add class namespace declared fields - for field_name, anno_type in namespace.get('__annotations__', {}).items(): - options = namespace.get(field_name, fields.FieldOptions()) + for field_name, annotated_type in namespace.get('__annotations__', {}).items(): + options = namespace.get(field_name, expressions.EmptyValue) if not isinstance(options, fields.FieldOptions): - # noinspection PyTypeChecker,SpellCheckingInspection - raise DocumentError( - loc=(name, field_name), - msg=f'Invalid field descriptor {type(options)}. ' - f'Use mongotoy.field() or mongotoy.reference() descriptors' - ) + options = fields.FieldOptions(default=options) + try: _fields[field_name] = fields.Field( - mapper=mappers.build_mapper(anno_type, options=options), + mapper=mappers.build_mapper(annotated_type, options=options), options=options ) except TypeError as e: # noinspection PyTypeChecker raise DocumentError( loc=(name, field_name), - msg=f'Invalid field annotation {anno_type}. {str(e)}' + msg=f'Invalid field annotation {annotated_type}. {str(e)}' ) from None + except Exception as e: # noinspection PyTypeChecker raise DocumentError( diff --git a/mongotoy/expressions.py b/mongotoy/expressions.py index e79d5dc..720bec5 100644 --- a/mongotoy/expressions.py +++ b/mongotoy/expressions.py @@ -1,8 +1,11 @@ import re +import typing from typing import Literal import pymongo +from mongotoy import geodata + EmptyValue = type('EmptyValue', (), {})() IndexType = Literal[-1, 1, '2d', '2dsphere', 'hashed', 'text'] @@ -258,6 +261,90 @@ def Regex(cls, field, value: re.Pattern) -> 'Query': """ return cls({str(field): {'$regex': value}}) + @classmethod + def Intersects(cls, field, value: geodata.Geometry) -> 'Query': + """ + Creates a geo intersects query expression. + + Selects documents whose geospatial data intersects with a specified GeoJSON object; + i.e. where the intersection of the data and the specified object is non-empty. + + Args: + field: The field name. + value (geodata.Geometry): The geometry. + + Returns: + Query: The geo intersects query expression. + """ + return cls({ + str(field): { + '$geoIntersects': { + '$geometry': value.dump_json() + } + } + }) + + @classmethod + def Within(cls, field, value: geodata.Polygon | geodata.MultiPolygon) -> 'Query': + """ + Creates a geo within query expression. + + Selects documents with geospatial data that exists entirely within a specified shape. + + Args: + field: The field name. + value (geodata.Polygon | geodata.MultiPolygon): The polygon or multipolygon geometry. + + Returns: + Query: The geo within query expression. + """ + return cls({ + str(field): { + '$geoWithin': { + '$geometry': value.dump_json() + } + } + }) + + @classmethod + def Near( + cls, + field, + value: geodata.Point, + max_distance: typing.Optional[int] = None, + min_distance: typing.Optional[int] = None, + as_near_sphere: bool = False + ) -> 'Query': + """ + Creates a geo near query expression. + + Specifies a point for which a geospatial query returns the documents from nearest to farthest + + Args: + field: The field name. + value (geodata.Point): The center point geometry. + max_distance (int, Optional): Max distance (in meters) from the center point. + min_distance (int, Optional): Min distance (in meters) from the center point. + as_near_sphere (bool): Calculates distances using spherical geometry. + + Returns: + Query: The geo near query expression. + """ + near_op = '$nearSphere' if as_near_sphere else '$near' + near_exp = { + '$geometry': value.dump_json() + } + if max_distance: + near_exp['$maxDistance'] = max_distance + if min_distance: + near_exp['$minDistance'] = min_distance + + return cls({ + str(field): { + near_op: near_exp + } + }) + # noinspection PyPep8Naming def Q(**kwargs) -> Query: diff --git a/mongotoy/geodata.py b/mongotoy/geodata.py index 111436a..a96b6b3 100644 --- a/mongotoy/geodata.py +++ b/mongotoy/geodata.py @@ -33,7 +33,8 @@ def dump_json(self) -> dict: Returns: dict: The GeoJSON representation of the geometry. """ - return {"type": self.type, "coordinates": self.coordinates} + extra = getattr(self, '__geo_extra__', {}) + return {"type": self.type, "coordinates": self.coordinates, **extra} def parse_geojson(geojson: dict, parser: typing.Type[Geometry]) -> Geometry: @@ -101,10 +102,11 @@ class Point(Position, Geometry): https://datatracker.ietf.org/doc/html/rfc7946#section-3.1.2 """ - def __init__(self, *coordinates: int | float): + def __init__(self, *coordinates: int | float, **kwargs): if len(coordinates) != 2: raise TypeError(f'The Point must represent single position, i.e. [lat, long]') super().__init__(*coordinates) + self.__geo_extra__ = kwargs class MultiPoint(list[Point], Geometry): @@ -115,8 +117,9 @@ class MultiPoint(list[Point], Geometry): https://datatracker.ietf.org/doc/html/rfc7946#section-3.1.3 """ - def __init__(self, *points: Point): + def __init__(self, *points: Point, **kwargs): super().__init__([Point(*i) for i in points]) + self.__geo_extra__ = kwargs class LineString(list[Point], Geometry): @@ -127,10 +130,11 @@ class LineString(list[Point], Geometry): https://datatracker.ietf.org/doc/html/rfc7946#section-3.1.4 """ - def __init__(self, *points: Point): + def __init__(self, *points: Point, **kwargs): if not len(points) >= 2: raise TypeError('The LineString must be an array of two or more Points') super().__init__([Point(*i) for i in points]) + self.__geo_extra__ = kwargs class MultiLineString(list[LineString], Geometry): @@ -141,8 +145,9 @@ class MultiLineString(list[LineString], Geometry): https://datatracker.ietf.org/doc/html/rfc7946#section-3.1.5 """ - def __init__(self, *lines: LineString): + def __init__(self, *lines: LineString, **kwargs): super().__init__([LineString(*i) for i in lines]) + self.__geo_extra__ = kwargs class Polygon(list[LineString], Geometry): @@ -150,8 +155,9 @@ class Polygon(list[LineString], Geometry): https://datatracker.ietf.org/doc/html/rfc7946#section-3.1.6 """ - def __init__(self, *rings: LineString): + def __init__(self, *rings: LineString, **kwargs): super().__init__([LinearRing(*i) for i in rings]) + self.__geo_extra__ = kwargs class MultiPolygon(list[Polygon], Geometry): @@ -162,5 +168,6 @@ class MultiPolygon(list[Polygon], Geometry): https://datatracker.ietf.org/doc/html/rfc7946#section-3.1.7 """ - def __init__(self, *polygons: Polygon): + def __init__(self, *polygons: Polygon, **kwargs): super().__init__([Polygon(*i) for i in polygons]) + self.__geo_extra__ = kwargs diff --git a/pyproject.toml b/pyproject.toml index 18ca9c2..8dd5d35 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "mongotoy" -version = "0.1.6" +version = "0.1.7" description = "Comprehensive ODM for MongoDB" license = "Apache-2.0" authors = ["gurcuff91 "]