extract_fields.py

from datetime import datetime
from hamilton.function_modifiers import extract_fields

def _fetch_dates__EMEA() -> dict:
   return {
       "a": "31/01/2000",
       "b": "28/02/2010",
       "c": "30/03/2020",
   }

def _fetch_dates__US() -> dict:
   return {
       "a": "01/31/2000",
       "b": "02/28/2010",
       "c": "03/30/2020",
   }

@extract_fields({"a": datetime, "b": datetime, "c": datetime})
def fetch_dates(region: str) -> dict:
   if region == "US":
       dates_US = _fetch_dates__US()
       return {
           k: datetime.strptime(
               dates_US[k],
               "%m/%d/%Y",
           )
           for k in dates_US
       }
   elif region == "EMEA":
       dates_EMEA = _fetch_dates__EMEA()
       return {
           k: datetime.strptime(
               dates_EMEA[k],
               "%d/%m/%Y",
           )
           for k in dates_EMEA
       }
   else:
       raise ValueError("Region unknown.")

def ab_period(a: datetime, b: datetime) -> int:
   return (b - a).days

def ac_period(a: datetime, c: datetime) -> int:
   return (c - a).days

def bc_period(b: datetime, c: datetime) -> int:
   return (c - b).days