path: root/dateparse/dateparse_utils.py

#!/usr/bin/env python3

import antlr4  # type: ignore
import datetime
import dateutil.easter
import holidays  # type: ignore
import re
import sys
from typing import Any, Dict, Optional

from dateparse.dateparse_utilsLexer import dateparse_utilsLexer  # type: ignore
from dateparse.dateparse_utilsListener import dateparse_utilsListener  # type: ignore
from dateparse.dateparse_utilsParser import dateparse_utilsParser  # type: ignore


class ParseException(Exception):
    def __init__(self, message: str) -> None:
        self.message = message


class DateParser(dateparse_utilsListener):
    PARSE_TYPE_SINGLE_DATE_EXPR = 1
    PARSE_TYPE_BASE_AND_OFFSET_EXPR = 2
    CONSTANT_DAYS = 7
    CONSTANT_WEEKS = 8
    CONSTANT_MONTHS = 9
    CONSTANT_YEARS = 10

    def __init__(self):
        self.month_name_to_number = {
            "jan": 1,
            "feb": 2,
            "mar": 3,
            "apr": 4,
            "may": 5,
            "jun": 6,
            "jul": 7,
            "aug": 8,
            "sep": 9,
            "oct": 10,
            "nov": 11,
            "dec": 12,
        }
        self.day_name_to_number = {
            "mon": 0,
            "tue": 1,
            "wed": 2,
            "thu": 3,
            "fri": 4,
            "sat": 5,
            "sun": 6,
        }
        self.delta_unit_to_constant = {
            "day": DateParser.CONSTANT_DAYS,
            "wee": DateParser.CONSTANT_WEEKS,
        }
        self.date: Optional[datetime.date] = None

    def parse_date_string(self, date_string: str) -> Optional[datetime.date]:
        input_stream = antlr4.InputStream(date_string)
        lexer = dateparse_utilsLexer(input_stream)
        stream = antlr4.CommonTokenStream(lexer)
        parser = dateparse_utilsParser(stream)
        tree = parser.parse()
        walker = antlr4.ParseTreeWalker()
        walker.walk(self, tree)
        return self.get_date()

    def get_date(self) -> Optional[datetime.date]:
        return self.date

    def enterDateExpr(self, ctx: dateparse_utilsParser.DateExprContext):
        self.date = None
        self.context: Dict[str, Any] = {}
        if ctx.singleDateExpr() is not None:
            self.main_type = DateParser.PARSE_TYPE_SINGLE_DATE_EXPR
        elif ctx.baseAndOffsetDateExpr() is not None:
            self.main_type = DateParser.PARSE_TYPE_BASE_AND_OFFSET_EXPR

    @staticmethod
    def normalize_special_day_name(name: str) -> str:
        name = name.lower()
        name = name.replace("'", "")
        name = name.replace("xmas", "christmas")
        name = name.replace("mlk", "martin luther king")
        name = name.replace(" ", "")
        eve = "eve" if name[-3:] == "eve" else ""
        name = name[:5] + eve
        name = name.replace("washi", "presi")
        return name

    def parse_special(self, name: str) -> Optional[datetime.date]:
        today = datetime.date.today()
        year = self.context.get("year", today.year)
        name = DateParser.normalize_special_day_name(self.context["special"])
        if name == "today":
            return today
        if name == "easte":
            return dateutil.easter.easter(year=year)
        for holiday_date, holiday_name in sorted(
            holidays.US(years=year).items()
        ):
            if "Observed" not in holiday_name:
                holiday_name = DateParser.normalize_special_day_name(
                    holiday_name
                )
                if name == holiday_name:
                    return holiday_date
        if name == "chriseve":
            return datetime.date(year=year, month=12, day=24)
        elif name == "newyeeve":
            return datetime.date(year=year, month=12, day=31)
        return None

    def parse_normal(self) -> datetime.date:
        if "month" not in self.context:
            raise ParseException("Missing month")
        if "day" not in self.context:
            raise ParseException("Missing day")
        if "year" not in self.context:
            today = datetime.date.today()
            self.context["year"] = today.year
        return datetime.date(
            year=int(self.context["year"]),
            month=int(self.context["month"]),
            day=int(self.context["day"]),
        )

    def exitDateExpr(self, ctx: dateparse_utilsParser.DateExprContext) -> None:
        """When we leave the date expression, populate self.date."""
        if "special" in self.context:
            self.date = self.parse_special(self.context["special"])
        else:
            self.date = self.parse_normal()
        assert self.date is not None

        # For a single date, just return the date we pulled out.
        if self.main_type == DateParser.PARSE_TYPE_SINGLE_DATE_EXPR:
            return

        # Otherwise treat self.date as a base date that we're modifying
        # with an offset.
        if not "delta_int" in self.context:
            raise ParseException("Missing delta_int?!")
        count = self.context["delta_int"]
        if count == 0:
            return

        # Adjust count's sign based on the presence of 'before' or 'after'.
        if "delta_before_after" in self.context:
            before_after = self.context["delta_before_after"].lower()
            if before_after == "before":
                count = -count

        # What are we counting units of?
        if "delta_unit" not in self.context:
            raise ParseException("Missing delta_unit?!")
        unit = self.context["delta_unit"]
        if unit == DateParser.CONSTANT_DAYS:
            timedelta = datetime.timedelta(days=count)
            self.date = self.date + timedelta
        elif unit == DateParser.CONSTANT_WEEKS:
            timedelta = datetime.timedelta(weeks=count)
            self.date = self.date + timedelta
        else:
            direction = 1 if count > 0 else -1
            count = abs(count)
            timedelta = datetime.timedelta(days=direction)

            while True:
                dow = self.date.weekday()
                if dow == unit:
                    count -= 1
                    if count == 0:
                        return
                self.date = self.date + timedelta

    def enterDeltaInt(self, ctx: dateparse_utilsParser.DeltaIntContext) -> None:
        try:
            i = int(ctx.getText())
        except:
            raise ParseException(f"Bad delta int: {ctx.getText()}")
        else:
            self.context["delta_int"] = i

    def enterDeltaUnit(
        self, ctx: dateparse_utilsParser.DeltaUnitContext
    ) -> None:
        try:
            txt = ctx.getText().lower()[:3]
            if txt in self.day_name_to_number:
                txt = self.day_name_to_number[txt]
            elif txt in self.delta_unit_to_constant:
                txt = self.delta_unit_to_constant[txt]
            else:
                raise ParseException(f"Bad delta unit: {ctx.getText()}")
        except:
            raise ParseException(f"Bad delta unit: {ctx.getText()}")
        else:
            self.context["delta_unit"] = txt

    def enterDeltaBeforeAfter(
        self, ctx: dateparse_utilsParser.DeltaBeforeAfterContext
    ) -> None:
        try:
            txt = ctx.getText().lower()
        except:
            raise ParseException(f"Bad delta before|after: {ctx.getText()}")
        else:
            self.context["delta_before_after"] = txt

    def exitNthWeekdayInMonthMaybeYearExpr(
        self, ctx: dateparse_utilsParser.NthWeekdayInMonthMaybeYearExprContext
    ) -> None:
        """Do a bunch of work to convert expressions like...

        'the 2nd Friday of June' -and-
        'the last Wednesday in October'

        ...into base + offset expressions instead.
        """
        try:
            if "nth" not in self.context:
                raise ParseException(f"Missing nth number: {ctx.getText()}")
            n = self.context["nth"]
            if n < 1 or n > 5:  # months never have more than 5 Foodays
                if n != -1:
                    raise ParseException(f"Invalid nth number: {ctx.getText()}")
            del self.context["nth"]
            self.context["delta_int"] = n

            year = self.context.get("year", datetime.date.today().year)
            if "month" not in self.context:
                raise ParseException(
                    f"Missing month expression: {ctx.getText()}"
                )
            month = self.context["month"]

            dow = self.context["dow"]
            del self.context["dow"]
            self.context["delta_unit"] = dow

            # For the nth Fooday in Month, start at the 1st of the
            # month and count ahead N Foodays.  For the last Fooday in
            # Month, start at the last of the month and count back one
            # Fooday.
            if n == -1:
                month += 1
                if month == 13:
                    month = 1
                    year += 1
                tmp_date = datetime.date(year=year, month=month, day=1)
                tmp_date = tmp_date - datetime.timedelta(days=1)

                self.context["year"] = tmp_date.year
                self.context["month"] = tmp_date.month
                self.context["day"] = tmp_date.day

                # The delta adjustment code can handle the case where
                # the last day of the month is the day we're looking
                # for already.
            else:
                self.context["year"] = year
                self.context["month"] = month
                self.context["day"] = 1
            self.main_type = DateParser.PARSE_TYPE_BASE_AND_OFFSET_EXPR
        except:
            raise ParseException(
                f"Invalid nthWeekday expression: {ctx.getText()}"
            )

    def exitFirstLastWeekdayInMonthMaybeYearExpr(
        self,
        ctx: dateparse_utilsParser.FirstLastWeekdayInMonthMaybeYearExprContext,
    ) -> None:
        self.exitNthWeekdayInMonthMaybeYearExpr(ctx)

    def enterNth(self, ctx: dateparse_utilsParser.NthContext) -> None:
        try:
            i = ctx.getText()
            m = re.match("\d+[a-z][a-z]", i)
            if m is not None:
                i = i[:-2]
            i = int(i)
        except:
            raise ParseException(f"Bad nth expression: {ctx.getText()}")
        else:
            self.context["nth"] = i

    def enterFirstOrLast(
        self, ctx: dateparse_utilsParser.FirstOrLastContext
    ) -> None:
        try:
            txt = ctx.getText()
            if txt == "first":
                txt = 1
            elif txt == "last":
                txt = -1
            else:
                raise ParseException(
                    f"Bad first|last expression: {ctx.getText()}"
                )
        except:
            raise ParseException(f"Bad first|last expression: {ctx.getText()}")
        else:
            self.context["nth"] = txt

    def enterDayName(self, ctx: dateparse_utilsParser.DayNameContext) -> None:
        try:
            dow = ctx.getText().lower()[:3]
            dow = self.day_name_to_number.get(dow, None)
        except:
            raise ParseException("Bad day of week")
        else:
            self.context["dow"] = dow

    def enterDayOfMonth(
        self, ctx: dateparse_utilsParser.DayOfMonthContext
    ) -> None:
        try:
            day = int(ctx.getText())
            if day < 1 or day > 31:
                raise ParseException(
                    f"Bad dayOfMonth expression: {ctx.getText()}"
                )
        except:
            raise ParseException(f"Bad dayOfMonth expression: {ctx.getText()}")
        self.context["day"] = day

    def enterMonthName(
        self, ctx: dateparse_utilsParser.MonthNameContext
    ) -> None:
        try:
            month = ctx.getText()
            month = month.lower()[:3]
            month = self.month_name_to_number.get(month, None)
            if month is None:
                raise ParseException(
                    f"Bad monthName expression: {ctx.getText()}"
                )
        except:
            raise ParseException(f"Bad monthName expression: {ctx.getText()}")
        else:
            self.context["month"] = month

    def enterMonthNumber(
        self, ctx: dateparse_utilsParser.MonthNumberContext
    ) -> None:
        try:
            month = int(ctx.getText())
            if month < 1 or month > 12:
                raise ParseException(
                    f"Bad monthNumber expression: {ctx.getText()}"
                )
        except:
            raise ParseException(f"Bad monthNumber expression: {ctx.getText()}")
        else:
            self.context["month"] = month

    def enterYear(self, ctx: dateparse_utilsParser.YearContext) -> None:
        try:
            year = int(ctx.getText())
            if year < 1:
                raise ParseException(f"Bad year expression: {ctx.getText()}")
        except:
            raise ParseException(f"Bad year expression: {ctx.getText()}")
        else:
            self.context["year"] = year

    def enterSpecialDate(
        self, ctx: dateparse_utilsParser.SpecialDateContext
    ) -> None:
        try:
            txt = ctx.getText().lower()
        except:
            raise ParseException(f"Bad specialDate expression: {ctx.getText()}")
        else:
            self.context["special"] = txt


def main() -> None:
    parser = DateParser()
    for line in sys.stdin:
        line = line.strip()
        line = line.lower()
        line = re.sub(r"#.*$", "", line)
        if re.match(r"^ *$", line) is not None:
            continue
        print(parser.parse_date_string(line))
    sys.exit(0)


if __name__ == "__main__":
    main()