generate_docstrings.py

"""This script ensures the Rust extension docstrings are identical to the
Python ones.

It does so by parsing the Python docstrings and generating a Rust file with the
same docstrings. This file is then included in the Rust extension.
"""

import enum
import inspect
import sys
from itertools import chain

from whenever import _pywhenever as W

assert sys.version_info >= (
    3,
    13,
), "This script requires Python 3.13 or later due to how docstrings are rendered."

classes = {
    cls
    for name, cls in W.__dict__.items()
    if (
        not name.startswith("_")
        and inspect.isclass(cls)
        and cls.__module__ == "whenever"
        and not issubclass(cls, enum.Enum)
    )
}
functions = {
    func
    for name, func in inspect.getmembers(W)
    if (
        not name.startswith("_")
        and inspect.isfunction(func)
        and func.__module__ == "whenever"
    )
}


methods = {
    getattr(cls, name)
    for cls in chain(
        classes,
        (
            # some methods are documented in their ABCs
            W._BasicConversions,
            W._KnowsLocal,
            W._KnowsInstant,
            W._KnowsInstantAndLocal,
        ),
    )
    for name, m in cls.__dict__.items()
    if (
        not name.startswith("_")
        and (
            inspect.isfunction(m)
            or
            # this catches classmethods
            inspect.ismethod(getattr(cls, name))
        )
    )
}

MAGIC_STRINGS = {
    (name, value)
    for name, value in W.__dict__.items()
    if isinstance(value, str) and name.isupper() and not name.startswith("_")
}

CSTR_TEMPLATE = 'pub(crate) const {varname}: &CStr = c"\\\n{doc}";'
STR_TEMPLATE = 'pub(crate) const {varname}: &str = "{value}";'
SIG_TEMPLATE = "{name}({self}, offset=0, /)\n--\n\n{doc}"
HEADER = """\
// Do not manually edit this file.
// It has been autogenerated by generate_docstrings.py
use std::ffi::CStr;
"""

MANUALLY_DEFINED_SIGS: dict[object, str] = {
    W.ZonedDateTime.add: """\
($self, delta=None, /, *, years=0, months=0, days=0, hours=0, \
minutes=0, seconds=0, milliseconds=0, microseconds=0, nanoseconds=0, \
disambiguate=None)""",
    W.ZonedDateTime.replace: """\
($self, /, *, year=None, month=None, day=None, hour=None, \
minute=None, second=None, nanosecond=None, tz=None, disambiguate)""",
    W.OffsetDateTime.add: """\
($self, delta=None, /, *, years=0, months=0, weeks=0, days=0, \
hours=0, minutes=0, seconds=0, milliseconds=0, microseconds=0, nanoseconds=0, \
ignore_dst=False)""",
    W.OffsetDateTime.replace: """\
($self, /, *, year=None, month=None, day=None, hour=None, \
minute=None, second=None, nanosecond=None, offset=None, ignore_dst=False)""",
    W.LocalDateTime.add: """\
($self, delta=None, /, *, years=0, months=0, days=0, \
hours=0, minutes=0, seconds=0, milliseconds=0, microseconds=0, nanoseconds=0, \
ignore_dst=False)""",
    W.LocalDateTime.replace: """\
($self, /, *, year=None, month=None, day=None, hour=None, \
minute=None, second=None, nanosecond=None)""",
    W.Date.replace: "($self, /, *, year=None, month=None, day=None)",
    W.MonthDay.replace: "($self, /, *, month=None, day=None)",
    W.Time.replace: "($self, /, *, hour=None, minute=None, second=None, nanosecond=None)",
    W.YearMonth.replace: "($self, /, *, year=None, month=None)",
    W.Instant.add: """\
($self, delta=None, /, *, hours=0, minutes=0, seconds=0, \
milliseconds=0, microseconds=0, nanoseconds=0)""",
}
MANUALLY_DEFINED_SIGS.update(
    {
        W.ZonedDateTime.subtract: MANUALLY_DEFINED_SIGS[W.ZonedDateTime.add],
        W.SystemDateTime.add: MANUALLY_DEFINED_SIGS[W.ZonedDateTime.add],
        W.SystemDateTime.subtract: MANUALLY_DEFINED_SIGS[W.ZonedDateTime.add],
        W.SystemDateTime.replace: MANUALLY_DEFINED_SIGS[
            W.ZonedDateTime.replace
        ],
        W.OffsetDateTime.subtract: MANUALLY_DEFINED_SIGS[W.OffsetDateTime.add],
        W.LocalDateTime.subtract: MANUALLY_DEFINED_SIGS[W.LocalDateTime.add],
        W.Instant.subtract: MANUALLY_DEFINED_SIGS[W.Instant.add],
    }
)
SKIP = {
    W._BasicConversions.format_common_iso,
    W._BasicConversions.from_py_datetime,
    W._BasicConversions.parse_common_iso,
    W._KnowsInstant.from_timestamp,
    W._KnowsInstant.from_timestamp_millis,
    W._KnowsInstant.from_timestamp_nanos,
    W._KnowsInstant.now,
    W._KnowsLocal.add,
    W._KnowsLocal.subtract,
    W._KnowsLocal.replace,
    W._KnowsLocal.replace_date,
    W._KnowsLocal.replace_time,
}


def method_doc(method):
    method.__annotations__.clear()
    try:
        sig = MANUALLY_DEFINED_SIGS[method]
    except KeyError:
        sig = (
            str(inspect.signature(method))
            # We use unicode escape of '(' to avoid messing up LSP in editors
            .replace("\u0028self", "\u0028$self").replace(
                "\u0028cls", "\u0028$type"
            )
        )
    doc = method.__doc__.replace('"', '\\"')
    return f"{method.__name__}{sig}\n--\n\n{doc}"


def print_everything():
    print(HEADER)
    for cls in sorted(classes, key=lambda x: x.__name__):
        assert cls.__doc__
        print(
            CSTR_TEMPLATE.format(
                varname=cls.__name__.upper(),
                doc=cls.__doc__.replace('"', '\\"'),
            )
        )

    for func in sorted(functions, key=lambda x: x.__name__):
        assert func.__doc__
        print(
            CSTR_TEMPLATE.format(
                varname=func.__name__.upper(),
                doc=func.__doc__.replace('"', '\\"'),
            )
        )

    for method in sorted(methods, key=lambda x: x.__qualname__):
        if method.__doc__ is None or method in SKIP:
            continue

        qualname = method.__qualname__
        if qualname.startswith("_"):
            qualname = qualname[1:]
        print(
            CSTR_TEMPLATE.format(
                varname=qualname.replace(".", "_").upper(),
                doc=method_doc(method),
            )
        )

    for name, value in sorted(MAGIC_STRINGS):
        print(STR_TEMPLATE.format(varname=name, value=value))


if __name__ == "__main__":
    print_everything()