bookwyrm/bookwyrm/utils/sealed_date.py

"""Implementation of the SealedDate class."""

from __future__ import annotations

from datetime import datetime, timedelta
import re
from typing import Any, Optional, Type, TypeVar, cast

from django.core.exceptions import ValidationError
from django.db import models
from django.forms import DateField
from django.forms.widgets import SelectDateWidget
from django.utils import timezone


__all__ = [
    "SealedDate",
    "from_partial_isoformat",
]

_partial_re = re.compile(r"(\d{4})(?:-(\d\d?))?(?:-(\d\d?))?$")
_westmost_tz = timezone.get_fixed_timezone(timedelta(hours=-12))

Sealed = TypeVar("Sealed", bound="SealedDate")  # TODO: use Self in Python >= 3.11

# TODO: migrate SealedDate: `datetime` => `date`
# TODO: migrate SealedDateField: `DateTimeField` => `DateField`


class SealedDate(datetime):
    """a date object sealed into a certain precision (day, month or year)"""

    @property
    def has_day(self) -> bool:
        return self.has_month

    @property
    def has_month(self) -> bool:
        return True

    def partial_isoformat(self) -> str:
        return self.strftime("%Y-%m-%d")

    @classmethod
    def from_datetime(cls: Type[Sealed], dt: datetime) -> Sealed:
        # pylint: disable=invalid-name
        if timezone.is_naive(dt):
            raise ValueError("naive datetime not accepted")
        return cls.combine(dt.date(), dt.time(), tzinfo=dt.tzinfo)

    @classmethod
    def from_date_parts(cls: Type[Sealed], year: int, month: int, day: int) -> Sealed:
        # because SealedDate is actually a datetime object, we must create it with a
        # timezone such that its date remains stable no matter the values of USE_TZ,
        # current_timezone and default_timezone.
        return cls.from_datetime(datetime(year, month, day, tzinfo=_westmost_tz))


class MonthSeal(SealedDate):
    @property
    def has_day(self) -> bool:
        return False

    def partial_isoformat(self) -> str:
        return self.strftime("%Y-%m")


class YearSeal(SealedDate):
    @property
    def has_month(self) -> bool:
        return False

    def partial_isoformat(self) -> str:
        return self.strftime("%Y")


def from_partial_isoformat(value: str) -> SealedDate:
    match = _partial_re.match(value)

    if not match:
        raise ValueError

    year, month, day = [val and int(val) for val in match.groups()]

    if month is None:
        return YearSeal.from_date_parts(year, 1, 1)
    elif day is None:
        return MonthSeal.from_date_parts(year, month, 1)
    else:
        return SealedDate.from_date_parts(year, month, day)


class SealedDateFormField(DateField):
    """date form field with support for SealedDate"""

    def prepare_value(self, value: Any) -> str:
        # As a convention, Django's `SelectDateWidget` uses "0" for missing
        # parts. We piggy-back into that, to make it work with SealedDate.
        if not isinstance(value, SealedDate):
            return cast(str, super().prepare_value(value))
        elif value.has_day:
            return value.strftime("%Y-%m-%d")
        elif value.has_month:
            return value.strftime("%Y-%m-0")
        else:
            return value.strftime("%Y-0-0")

    def to_python(self, value: Any) -> Optional[SealedDate]:
        try:
            date = super().to_python(value)
        except ValidationError as ex:
            if match := SelectDateWidget.date_re.match(value):
                year, month, day = map(int, match.groups())
            if not match or (day and not month) or not year:
                raise ex from None
            if not month:
                return YearSeal.from_date_parts(year, 1, 1)
            elif not day:
                return MonthSeal.from_date_parts(year, month, 1)
        else:
            if date is None:
                return None
            else:
                year, month, day = date.year, date.month, date.day

        return SealedDate.from_date_parts(year, month, day)


class SealedDateDescriptor:

    _SEAL_TYPES = {
        YearSeal: "YEAR",
        MonthSeal: "MONTH",
        SealedDate: "DAY",
    }

    _DATE_CLASSES = {
        "YEAR": YearSeal,
        "MONTH": MonthSeal,
    }

    def __init__(self, field):
        self.field = field

    @property
    def precision_field(self):
        """the name of the accompanying precision field"""
        return self.make_precision_name(self.field.attname)

    @classmethod
    def make_precision_name(cls, date_attr_name):
        # used by SealedDateField to make the name from the outside.
        # TODO: migrate to an attribute there?
        return f"{date_attr_name}_precision"

    @property
    def precision_choices(self):
        return (("DAY", "Day seal"), ("MONTH", "Month seal"), ("YEAR", "Year seal"))

    def __get__(self, instance, cls=None):
        if instance is None:
            return self

        value = instance.__dict__.get(self.field.attname)

        if not value or isinstance(value, SealedDate):
            return value

        # use precision field to construct SealedDate.
        seal_type = getattr(instance, self.precision_field, None)
        date_class = self._DATE_CLASSES.get(seal_type, SealedDate)

        return date_class.from_datetime(value)  # FIXME: drop datetimes.

    def __set__(self, instance, value):
        """assign value, with precision where available"""
        try:
            seal_type = self._SEAL_TYPES[value.__class__]
        except KeyError:
            value = self.field.to_python(value)
        else:
            setattr(instance, self.precision_field, seal_type)

        instance.__dict__[self.field.attname] = value


class SealedDateField(models.DateTimeField):  # FIXME: use DateField.

    descriptor_class = SealedDateDescriptor

    def formfield(self, **kwargs):
        kwargs.setdefault("form_class", SealedDateFormField)
        return super().formfield(**kwargs)

    # pylint: disable-next=arguments-renamed
    def contribute_to_class(self, model, our_name_in_model, **kwargs):
        # Define precision field.
        descriptor = self.descriptor_class(self)
        precision = models.CharField(
            null=True,
            blank=True,
            editable=False,
            max_length=10,
            choices=descriptor.precision_choices,
        )
        precision_name = descriptor.make_precision_name(our_name_in_model)

        model.add_to_class(precision_name, precision)
        return super().contribute_to_class(model, our_name_in_model, **kwargs)