Work-in-progress commit for type annotations in line item discounter

boxoffice/models/discount_policy.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence
 from uuid import UUID
 import secrets
 import string
@@ -203,8 +204,8 @@ def make_bulk(cls, discount_code_base, **kwargs) -> DiscountPolicy:
 
     @classmethod
     def get_from_ticket(
-        cls, ticket: Item, qty, coupon_codes=()
-    ) -> List[Tuple[DiscountPolicy, Optional[DiscountCoupon]]]:
+        cls, ticket: Item, qty, coupon_codes: Sequence[str] = ()
+    ) -> List[PolicyCoupon]:
         """
         Return a list of (discount_policy, discount_coupon) tuples.
 
@@ -214,9 +215,7 @@ def get_from_ticket(
             cls.discount_type == DiscountTypeEnum.AUTOMATIC,
             cls.item_quantity_min <= qty,
         ).all()
-        policies: List[Tuple[DiscountPolicy, Optional[DiscountCoupon]]] = [
-            (discount, None) for discount in automatic_discounts
-        ]
+        policies = [PolicyCoupon(discount, None) for discount in automatic_discounts]
         if not coupon_codes:
             return policies
         coupon_policies = ticket.discount_policies.filter(
@@ -250,7 +249,7 @@ def get_from_ticket(
 
             for coupon in coupons:
                 if coupon.usage_limit > coupon.used_count:
-                    policies.append((coupon.discount_policy, coupon))
+                    policies.append(PolicyCoupon(coupon.discount_policy, coupon))
         return policies
 
     @property
@@ -348,6 +347,12 @@ def update_used_count(self):
         )
 
 
+@dataclass
+class PolicyCoupon:
+    policy: DiscountPolicy
+    coupon: Optional[DiscountCoupon]
+
+
 create_title_trgm_trigger = sa.DDL(
     'CREATE INDEX idx_discount_policy_title_trgm on discount_policy'
     ' USING gin (title gin_trgm_ops);'

boxoffice/models/line_item.py

@@ -1,23 +1,23 @@
 from __future__ import annotations
 
-from collections import OrderedDict, namedtuple
+from collections import OrderedDict
 from datetime import date, datetime, timedelta, tzinfo
 from decimal import Decimal
 from typing import (
     TYPE_CHECKING,
     Dict,
-    Iterable,
+    List,
+    NamedTuple,
     Optional,
     Sequence,
     Union,
     cast,
-    overload,
 )
 from uuid import UUID
 
 from flask import current_app
 from isoweek import Week
-from typing_extensions import Literal
+from typing_extensions import TypedDict
 
 from baseframe import localize_timezone
 from coaster.utils import isoweek_datetime, midnight_to_utc, utcnow
@@ -38,30 +38,20 @@
 __all__ = ['LineItem', 'Assignee']
 
 
-LineItemTuple = namedtuple(
-    'LineItemTuple',
-    [
-        'ticket_id',
-        'id',
-        'base_amount',
-        'discount_policy_id',
-        'discount_coupon_id',
-        'discounted_amount',
-        'final_amount',
-    ],
-)
+class LineItemTuple(NamedTuple):
+    """Duck-type for LineItem."""
 
+    id: Optional[UUID]  # noqa: A003
+    ticket_id: UUID
+    base_amount: Optional[Decimal]
+    discount_policy_id: Optional[UUID] = None
+    discount_coupon_id: Optional[UUID] = None
+    discounted_amount: Optional[Decimal] = Decimal(0)
+    final_amount: Optional[Decimal] = None
 
-def make_ntuple(ticket_id, base_amount, **kwargs):
-    return LineItemTuple(
-        ticket_id,
-        kwargs.get('line_item_id', None),
-        base_amount,
-        kwargs.get('discount_policy_id', None),
-        kwargs.get('discount_coupon_id', None),
-        kwargs.get('discounted_amount', Decimal(0)),
-        kwargs.get('final_amount', None),
-    )
+
+class LineItemDict(TypedDict):
+    ticket_id: str
 
 
 class Assignee(BaseMixin, Model):
@@ -151,77 +141,51 @@ def permissions(self, actor, inherited=None):
             perms.add('org_admin')
         return perms
 
-    @overload
-    @classmethod
-    def calculate(
-        cls,
-        line_items: Sequence[LineItem],
-        realculate: Literal[True],
-        coupons: Optional[Sequence[str]] = None,
-    ):
-        ...
-
-    @overload
     @classmethod
     def calculate(
         cls,
-        line_items: Sequence[dict],
-        realculate: Literal[False],
+        line_items: Sequence[Union[LineItem, LineItemDict]],
         coupons: Optional[Sequence[str]] = None,
-    ):
-        ...
-
-    # TODO: Fix this classmethod's typing
-    @classmethod
-    def calculate(
-        cls,
-        line_items: Iterable[Union[LineItem, dict]],  # FIXME
-        recalculate=False,
-        coupons=None,
-    ):
+    ) -> Sequence[LineItemTuple]:
         """
         Return line item data tuples.
 
         For each line item, returns a tuple of base_amount, discounted_amount,
         final_amount, discount_policy and discount coupon populated
-
-        If the `recalculate` flag is set to `True`, the line_items will be considered
-        as SQLAlchemy objects.
         """
-        item_line_items = {}
-        calculated_line_items = []
+        base_amount: Optional[Decimal]
+        item_line_items: Dict[str, List[LineItemTuple]] = {}
+        calculated_line_items: List[LineItemTuple] = []
         coupon_list = list(set(coupons)) if coupons else []
         discounter = LineItemDiscounter()
 
         # make named tuples for line items,
         # assign the base_amount for each of them, None if a ticket is unavailable
         for line_item in line_items:
-            if recalculate:
+            ticket: Optional[Item]
+            if isinstance(line_item, LineItem):
                 ticket = line_item.ticket
-                # existing line item, use the original base amount
                 base_amount = line_item.base_amount
-                line_item_id = line_item.id
+                line_item_id = cast(UUID, line_item.id)
             else:
                 ticket = Item.query.get(line_item['ticket_id'])
-                # new line item, use the current price
                 if ticket is None:
-                    base_amount = None
-                else:
-                    current_price = ticket.current_price()
-                    base_amount = (
-                        current_price.amount
-                        if current_price is not None and ticket.is_available
-                        else None
-                    )
+                    continue
+                current_price = ticket.current_price()
+                base_amount = (
+                    current_price.amount
+                    if current_price is not None and ticket.is_available
+                    else None
+                )
                 line_item_id = None
 
             if not item_line_items.get(str(ticket.id)):
                 item_line_items[str(ticket.id)] = []
             item_line_items[str(ticket.id)].append(
-                make_ntuple(
-                    ticket_id=ticket.id,
+                LineItemTuple(
+                    id=line_item_id,
+                    ticket_id=cast(UUID, ticket.id),
                     base_amount=base_amount,
-                    line_item_id=line_item_id,
                 )
             )