plain.postgres 0.84.0__py3-none-any.whl

plain/postgres/functions/datetime.py

@@ -0,0 +1,454 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from plain.postgres.dialect import (
+    date_extract_sql,
+    date_trunc_sql,
+    datetime_cast_date_sql,
+    datetime_cast_time_sql,
+    datetime_extract_sql,
+    datetime_trunc_sql,
+    time_extract_sql,
+    time_trunc_sql,
+)
+from plain.postgres.expressions import Func
+from plain.postgres.fields import (
+    DateField,
+    DateTimeField,
+    DurationField,
+    Field,
+    IntegerField,
+    TimeField,
+)
+from plain.postgres.lookups import (
+    Transform,
+    YearExact,
+    YearGt,
+    YearGte,
+    YearLt,
+    YearLte,
+)
+from plain.utils import timezone
+
+if TYPE_CHECKING:
+    from plain.postgres.connection import DatabaseConnection
+    from plain.postgres.sql.compiler import SQLCompiler
+
+
+class TimezoneMixin(Transform):
+    tzinfo = None
+
+    def get_tzname(self) -> str | None:
+        # Timezone conversions must happen to the input datetime *before*
+        # applying a function. 2015-12-31 23:00:00 -02:00 is stored in the
+        # database as 2016-01-01 01:00:00 +00:00. Any results should be
+        # based on the input datetime not the stored datetime.
+        if self.tzinfo is None:
+            return timezone.get_current_timezone_name()
+        else:
+            return timezone._get_timezone_name(self.tzinfo)
+
+
+class Extract(TimezoneMixin, Transform):
+    lookup_name: str | None = None
+    output_field = IntegerField()
+
+    def __init__(
+        self,
+        expression: Any,
+        lookup_name: str | None = None,
+        tzinfo: Any = None,
+        **extra: Any,
+    ) -> None:
+        if self.lookup_name is None:
+            self.lookup_name = lookup_name
+        if self.lookup_name is None:
+            raise ValueError("lookup_name must be provided")
+        self.tzinfo = tzinfo
+        super().__init__(expression, **extra)
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        # lookup_name is guaranteed to be str after __init__ validation
+        assert self.lookup_name is not None
+        sql, params = compiler.compile(self.lhs)
+        lhs_output_field = self.lhs.output_field
+        if isinstance(lhs_output_field, DateTimeField):
+            tzname = self.get_tzname()
+            sql, params = datetime_extract_sql(
+                self.lookup_name, sql, tuple(params), tzname
+            )
+        elif self.tzinfo is not None:
+            raise ValueError("tzinfo can only be used with DateTimeField.")
+        elif isinstance(lhs_output_field, DateField):
+            sql, params = date_extract_sql(self.lookup_name, sql, tuple(params))
+        elif isinstance(lhs_output_field, TimeField):
+            sql, params = time_extract_sql(self.lookup_name, sql, tuple(params))
+        elif isinstance(lhs_output_field, DurationField):
+            # PostgreSQL has native duration (interval) type
+            sql, params = time_extract_sql(self.lookup_name, sql, tuple(params))
+        else:
+            # resolve_expression has already validated the output_field so this
+            # assert should never be hit.
+            raise ValueError("Tried to Extract from an invalid type.")
+        return sql, list(params)
+
+    def resolve_expression(
+        self,
+        query: Any = None,
+        allow_joins: bool = True,
+        reuse: Any = None,
+        summarize: bool = False,
+        for_save: bool = False,
+    ) -> Extract:
+        copy = super().resolve_expression(
+            query, allow_joins, reuse, summarize, for_save
+        )
+        field = getattr(copy.lhs, "output_field", None)
+        if field is None:
+            return copy
+        if not isinstance(field, DateField | DateTimeField | TimeField | DurationField):
+            raise ValueError(
+                "Extract input expression must be DateField, DateTimeField, "
+                "TimeField, or DurationField."
+            )
+        # Passing dates to functions expecting datetimes is most likely a mistake.
+        if type(field) == DateField and copy.lookup_name in (  # noqa: E721
+            "hour",
+            "minute",
+            "second",
+        ):
+            raise ValueError(
+                f"Cannot extract time component '{copy.lookup_name}' from DateField '{field.name}'."
+            )
+        if isinstance(field, DurationField) and copy.lookup_name in (
+            "year",
+            "iso_year",
+            "month",
+            "week",
+            "week_day",
+            "iso_week_day",
+            "quarter",
+        ):
+            raise ValueError(
+                f"Cannot extract component '{copy.lookup_name}' from DurationField '{field.name}'."
+            )
+        return copy
+
+
+class ExtractYear(Extract):
+    lookup_name = "year"
+
+
+class ExtractIsoYear(Extract):
+    """Return the ISO-8601 week-numbering year."""
+
+    lookup_name = "iso_year"
+
+
+class ExtractMonth(Extract):
+    lookup_name = "month"
+
+
+class ExtractDay(Extract):
+    lookup_name = "day"
+
+
+class ExtractWeek(Extract):
+    """
+    Return 1-52 or 53, based on ISO-8601, i.e., Monday is the first of the
+    week.
+    """
+
+    lookup_name = "week"
+
+
+class ExtractWeekDay(Extract):
+    """
+    Return Sunday=1 through Saturday=7.
+
+    To replicate this in Python: (mydatetime.isoweekday() % 7) + 1
+    """
+
+    lookup_name = "week_day"
+
+
+class ExtractIsoWeekDay(Extract):
+    """Return Monday=1 through Sunday=7, based on ISO-8601."""
+
+    lookup_name = "iso_week_day"
+
+
+class ExtractQuarter(Extract):
+    lookup_name = "quarter"
+
+
+class ExtractHour(Extract):
+    lookup_name = "hour"
+
+
+class ExtractMinute(Extract):
+    lookup_name = "minute"
+
+
+class ExtractSecond(Extract):
+    lookup_name = "second"
+
+
+DateField.register_lookup(ExtractYear)
+DateField.register_lookup(ExtractMonth)
+DateField.register_lookup(ExtractDay)
+DateField.register_lookup(ExtractWeekDay)
+DateField.register_lookup(ExtractIsoWeekDay)
+DateField.register_lookup(ExtractWeek)
+DateField.register_lookup(ExtractIsoYear)
+DateField.register_lookup(ExtractQuarter)
+
+TimeField.register_lookup(ExtractHour)
+TimeField.register_lookup(ExtractMinute)
+TimeField.register_lookup(ExtractSecond)
+
+DateTimeField.register_lookup(ExtractHour)
+DateTimeField.register_lookup(ExtractMinute)
+DateTimeField.register_lookup(ExtractSecond)
+
+ExtractYear.register_lookup(YearExact)
+ExtractYear.register_lookup(YearGt)
+ExtractYear.register_lookup(YearGte)
+ExtractYear.register_lookup(YearLt)
+ExtractYear.register_lookup(YearLte)
+
+ExtractIsoYear.register_lookup(YearExact)
+ExtractIsoYear.register_lookup(YearGt)
+ExtractIsoYear.register_lookup(YearGte)
+ExtractIsoYear.register_lookup(YearLt)
+ExtractIsoYear.register_lookup(YearLte)
+
+
+class Now(Func):
+    # STATEMENT_TIMESTAMP() returns the time at the start of the current statement,
+    # as opposed to CURRENT_TIMESTAMP which returns the time at the start of the
+    # transaction.
+    template = "STATEMENT_TIMESTAMP()"
+    output_field = DateTimeField()
+
+
+class TruncBase(TimezoneMixin, Transform):
+    kind: str | None = None
+
+    def __init__(
+        self,
+        expression: Any,
+        output_field: Field | None = None,
+        tzinfo: Any = None,
+        **extra: Any,
+    ) -> None:
+        self.tzinfo = tzinfo
+        super().__init__(expression, output_field=output_field, **extra)
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        # kind is guaranteed to be str in subclasses
+        assert self.kind is not None
+        sql, params = compiler.compile(self.lhs)
+        tzname = None
+        if isinstance(self.lhs.output_field, DateTimeField):
+            tzname = self.get_tzname()
+        elif self.tzinfo is not None:
+            raise ValueError("tzinfo can only be used with DateTimeField.")
+        if isinstance(self.output_field, DateTimeField):
+            sql, params = datetime_trunc_sql(self.kind, sql, tuple(params), tzname)
+        elif isinstance(self.output_field, DateField):
+            sql, params = date_trunc_sql(self.kind, sql, tuple(params), tzname)
+        elif isinstance(self.output_field, TimeField):
+            sql, params = time_trunc_sql(self.kind, sql, tuple(params), tzname)
+        else:
+            raise ValueError(
+                "Trunc only valid on DateField, TimeField, or DateTimeField."
+            )
+        return sql, list(params)
+
+    def resolve_expression(
+        self,
+        query: Any = None,
+        allow_joins: bool = True,
+        reuse: Any = None,
+        summarize: bool = False,
+        for_save: bool = False,
+    ) -> TruncBase:
+        copy = super().resolve_expression(
+            query, allow_joins, reuse, summarize, for_save
+        )
+        field = copy.lhs.output_field
+        # DateTimeField is a subclass of DateField so this works for both.
+        if not isinstance(field, DateField | TimeField):
+            raise TypeError(
+                f"{field.name!r} isn't a DateField, TimeField, or DateTimeField."
+            )
+        # If self.output_field was None, then accessing the field will trigger
+        # the resolver to assign it to self.lhs.output_field.
+        if not isinstance(copy.output_field, DateField | DateTimeField | TimeField):
+            raise ValueError(
+                "output_field must be either DateField, TimeField, or DateTimeField"
+            )
+        # Passing dates or times to functions expecting datetimes is most
+        # likely a mistake.
+        class_output_field = (
+            self.__class__.output_field
+            if isinstance(self.__class__.output_field, Field)
+            else None
+        )
+        output_field = class_output_field or copy.output_field
+        has_explicit_output_field = (
+            class_output_field or field.__class__ is not copy.output_field.__class__
+        )
+        if type(field) == DateField and (  # noqa: E721
+            isinstance(output_field, DateTimeField)
+            or copy.kind in ("hour", "minute", "second", "time")
+        ):
+            raise ValueError(
+                "Cannot truncate DateField '{}' to {}.".format(
+                    field.name,
+                    output_field.__class__.__name__
+                    if has_explicit_output_field
+                    else "DateTimeField",
+                )
+            )
+        elif isinstance(field, TimeField) and (
+            isinstance(output_field, DateTimeField)
+            or copy.kind in ("year", "quarter", "month", "week", "day", "date")
+        ):
+            raise ValueError(
+                "Cannot truncate TimeField '{}' to {}.".format(
+                    field.name,
+                    output_field.__class__.__name__
+                    if has_explicit_output_field
+                    else "DateTimeField",
+                )
+            )
+        return copy
+
+    def convert_value(
+        self, value: Any, expression: Any, connection: DatabaseConnection
+    ) -> Any:
+        if isinstance(self.output_field, DateTimeField):
+            if value is not None:
+                value = value.replace(tzinfo=None)
+                value = timezone.make_aware(value, self.tzinfo)
+        elif isinstance(value, datetime):
+            if value is None:
+                pass
+            elif isinstance(self.output_field, DateField):
+                value = value.date()
+            elif isinstance(self.output_field, TimeField):
+                value = value.time()
+        return value
+
+
+class Trunc(TruncBase):
+    def __init__(
+        self,
+        expression: Any,
+        kind: str,
+        output_field: Field | None = None,
+        tzinfo: Any = None,
+        **extra: Any,
+    ) -> None:
+        self.kind = kind
+        super().__init__(expression, output_field=output_field, tzinfo=tzinfo, **extra)
+
+
+class TruncYear(TruncBase):
+    kind = "year"
+
+
+class TruncQuarter(TruncBase):
+    kind = "quarter"
+
+
+class TruncMonth(TruncBase):
+    kind = "month"
+
+
+class TruncWeek(TruncBase):
+    """Truncate to midnight on the Monday of the week."""
+
+    kind = "week"
+
+
+class TruncDay(TruncBase):
+    kind = "day"
+
+
+class TruncDate(TruncBase):
+    kind = "date"
+    lookup_name = "date"
+    output_field = DateField()
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        # Cast to date rather than truncate to date.
+        sql, params = compiler.compile(self.lhs)
+        tzname = self.get_tzname()
+        sql, params = datetime_cast_date_sql(sql, tuple(params), tzname)
+        return sql, list(params)
+
+
+class TruncTime(TruncBase):
+    kind = "time"
+    lookup_name = "time"
+    output_field = TimeField()
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        # Cast to time rather than truncate to time.
+        sql, params = compiler.compile(self.lhs)
+        tzname = self.get_tzname()
+        sql, params = datetime_cast_time_sql(sql, tuple(params), tzname)
+        return sql, list(params)
+
+
+class TruncHour(TruncBase):
+    kind = "hour"
+
+
+class TruncMinute(TruncBase):
+    kind = "minute"
+
+
+class TruncSecond(TruncBase):
+    kind = "second"
+
+
+DateTimeField.register_lookup(TruncDate)
+DateTimeField.register_lookup(TruncTime)

plain/postgres/functions/math.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from plain.postgres.expressions import Func
+from plain.postgres.fields import Field
+from plain.postgres.functions.mixins import (
+    FixDecimalInputMixin,
+    NumericOutputFieldMixin,
+)
+from plain.postgres.lookups import Transform
+
+if TYPE_CHECKING:
+    pass
+
+
+class Abs(Transform):
+    function = "ABS"
+    lookup_name = "abs"
+
+
+class ACos(NumericOutputFieldMixin, Transform):
+    function = "ACOS"
+    lookup_name = "acos"
+
+
+class ASin(NumericOutputFieldMixin, Transform):
+    function = "ASIN"
+    lookup_name = "asin"
+
+
+class ATan(NumericOutputFieldMixin, Transform):
+    function = "ATAN"
+    lookup_name = "atan"
+
+
+class ATan2(NumericOutputFieldMixin, Func):
+    function = "ATAN2"
+    arity = 2
+
+
+class Ceil(Transform):
+    function = "CEILING"
+    lookup_name = "ceil"
+
+
+class Cos(NumericOutputFieldMixin, Transform):
+    function = "COS"
+    lookup_name = "cos"
+
+
+class Cot(NumericOutputFieldMixin, Transform):
+    function = "COT"
+    lookup_name = "cot"
+
+
+class Degrees(NumericOutputFieldMixin, Transform):
+    function = "DEGREES"
+    lookup_name = "degrees"
+
+
+class Exp(NumericOutputFieldMixin, Transform):
+    function = "EXP"
+    lookup_name = "exp"
+
+
+class Floor(Transform):
+    function = "FLOOR"
+    lookup_name = "floor"
+
+
+class Ln(NumericOutputFieldMixin, Transform):
+    function = "LN"
+    lookup_name = "ln"
+
+
+class Log(FixDecimalInputMixin, NumericOutputFieldMixin, Func):
+    function = "LOG"
+    arity = 2
+
+
+class Mod(FixDecimalInputMixin, NumericOutputFieldMixin, Func):
+    function = "MOD"
+    arity = 2
+
+
+class Pi(NumericOutputFieldMixin, Func):
+    function = "PI"
+    arity = 0
+
+
+class Power(NumericOutputFieldMixin, Func):
+    function = "POWER"
+    arity = 2
+
+
+class Radians(NumericOutputFieldMixin, Transform):
+    function = "RADIANS"
+    lookup_name = "radians"
+
+
+class Random(NumericOutputFieldMixin, Func):
+    function = "RANDOM"
+    arity = 0
+
+    def get_group_by_cols(self) -> list[Any]:
+        return []
+
+
+class Round(FixDecimalInputMixin, Transform):
+    function = "ROUND"
+    lookup_name = "round"
+    arity = None  # Override Transform's arity=1 to enable passing precision.
+
+    def __init__(self, expression: Any, precision: int = 0, **extra: Any) -> None:
+        super().__init__(expression, precision, **extra)
+
+    def _resolve_output_field(self) -> Field:
+        source = self.get_source_expressions()[0]
+        return source.output_field
+
+
+class Sign(Transform):
+    function = "SIGN"
+    lookup_name = "sign"
+
+
+class Sin(NumericOutputFieldMixin, Transform):
+    function = "SIN"
+    lookup_name = "sin"
+
+
+class Sqrt(NumericOutputFieldMixin, Transform):
+    function = "SQRT"
+    lookup_name = "sqrt"
+
+
+class Tan(NumericOutputFieldMixin, Transform):
+    function = "TAN"
+    lookup_name = "tan"

plain/postgres/functions/mixins.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, Any
+
+from plain.postgres.expressions import Func
+from plain.postgres.fields import DecimalField, Field, FloatField, IntegerField
+from plain.postgres.functions import Cast
+
+if TYPE_CHECKING:
+    from plain.postgres.connection import DatabaseConnection
+    from plain.postgres.sql.compiler import SQLCompiler
+
+
+class FixDecimalInputMixin(Func):
+    """
+    Mixin for Func subclasses that need to convert FloatField to DecimalField.
+
+    PostgreSQL doesn't support the following function signatures:
+    - LOG(double, double)
+    - MOD(double, double)
+    """
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: DatabaseConnection,
+        function: str | None = None,
+        template: str | None = None,
+        arg_joiner: str | None = None,
+        **extra_context: Any,
+    ) -> tuple[str, list[Any]]:
+        output_field = DecimalField(decimal_places=sys.float_info.dig, max_digits=1000)
+
+        clone = self.copy()
+        clone.set_source_expressions(
+            [
+                Cast(expression, output_field)
+                if isinstance(expression.output_field, FloatField)
+                else expression
+                for expression in self.get_source_expressions()
+            ]
+        )
+        return super(FixDecimalInputMixin, clone).as_sql(
+            compiler, connection, **extra_context
+        )
+
+
+class NumericOutputFieldMixin(Func):
+    def _resolve_output_field(self) -> DecimalField | FloatField | Field:
+        source_fields = self.get_source_fields()
+        if any(isinstance(s, DecimalField) for s in source_fields):
+            return DecimalField()
+        if any(isinstance(s, IntegerField) for s in source_fields):
+            return FloatField()
+        if source_fields:
+            if result := super()._resolve_output_field():
+                return result
+        return FloatField()