bookwyrm/bookwyrm/utils/partial_date.py

"""Implementation of the PartialDate 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

# pylint: disable=no-else-return

__all__ = [
    "PartialDate",
    "PartialDateModel",
    "from_partial_isoformat",
]

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

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

# TODO: migrate PartialDate: `datetime` => `date`
# TODO: migrate PartialDateModel: `DateTimeField` => `DateField`


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

    @property
    def has_day(self) -> bool:
        """whether this is a full date"""
        return self.has_month

    @property
    def has_month(self) -> bool:
        """whether this date includes month"""
        return True

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

    @classmethod
    def from_datetime(cls: Type[Partial], dt: datetime) -> Partial:
        """construct a PartialDate object from a timezone-aware datetime

        Use subclasses to specify precision. If `dt` is naive, `ValueError`
        is raised.
        """
        # 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[Partial], year: int, month: int, day: int) -> Partial:
        """construct a PartialDate from year, month, day.

        Use sublcasses to specify precision."""
        # because PartialDate 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 MonthParts(PartialDate):
    """a date bound into month precision"""

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

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


class YearParts(PartialDate):
    """a date bound into year precision"""

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

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


def from_partial_isoformat(value: str) -> PartialDate:
    """construct PartialDate from a partial string.

    Accepted formats: YYYY, YYYY-MM, YYYY-MM-DD; otherwise `ValueError`
    is raised.
    """
    match = _partial_re.match(value)

    if not match:
        raise ValueError

    year, month, day = [int(val) if val else -1 for val in match.groups()]

    if month < 0:
        return YearParts.from_date_parts(year, 1, 1)
    elif day < 0:
        return MonthParts.from_date_parts(year, month, 1)
    else:
        return PartialDate.from_date_parts(year, month, day)


class PartialDateFormField(DateField):
    """date form field with support for PartialDate"""

    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 PartialDate.
        if not isinstance(value, PartialDate):
            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[PartialDate]:
        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 YearParts.from_date_parts(year, 1, 1)
            elif not day:
                return MonthParts.from_date_parts(year, month, 1)
        else:
            if date is None:
                return None
            else:
                year, month, day = date.year, date.month, date.day

        return PartialDate.from_date_parts(year, month, day)


# For typing field and descriptor, below.
_SetType = datetime
_GetType = Optional[PartialDate]


class PartialDateDescriptor:
    """descriptor for PartialDateModel.

    Encapsulates the "two columns, one field" for PartialDateModel.
    """

    _PRECISION_NAMES: dict[Type[_SetType], str] = {
        YearParts: "YEAR",
        MonthParts: "MONTH",
        PartialDate: "DAY",
    }

    _PARTIAL_CLASSES: dict[Any, Type[PartialDate]] = {
        "YEAR": YearParts,
        "MONTH": MonthParts,
    }

    def __init__(self, field: models.Field[_SetType, _GetType]):
        self.field = field

    def __get__(self, instance: models.Model, cls: Any = None) -> _GetType:
        if instance is None:
            return self

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

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

        # use precision field to construct PartialDate.
        precision = getattr(instance, self.precision_field, None)
        date_class = self._PARTIAL_CLASSES.get(precision, PartialDate)

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

    def __set__(self, instance: models.Model, value: _SetType) -> None:
        """assign value, with precision where available"""
        try:
            precision = self._PRECISION_NAMES[value.__class__]
        except KeyError:
            value = self.field.to_python(value)
        else:
            setattr(instance, self.precision_field, precision)

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

    @classmethod
    def make_precision_name(cls, date_attr_name: str) -> str:
        """derive the precision field name from main attr name"""
        return f"{date_attr_name}_precision"

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

    @property
    def precision_choices(self) -> list[tuple[str, str]]:
        """valid options for precision database field"""
        return [("DAY", "Day prec."), ("MONTH", "Month prec."), ("YEAR", "Year prec.")]


class PartialDateModel(models.DateTimeField):  # type: ignore
    """a date field for Django models, using PartialDate as values"""

    descriptor_class = PartialDateDescriptor

    def formfield(self, **kwargs):  # type: ignore
        kwargs.setdefault("form_class", PartialDateFormField)
        return super().formfield(**kwargs)

    # pylint: disable-next=arguments-renamed
    def contribute_to_class(self, model, our_name_in_model, **kwargs):  # type: ignore
        # Define precision field.
        descriptor = self.descriptor_class(self)
        precision: models.Field[Optional[str], Optional[str]] = 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)