forerad

Utilities for collecting and analyzing with Citibike data in Python
Log | Files | Refs | README

utils.py (989B)

      1 import re
      2 import logging
      3 import datetime
      4 import pytz
      5 
      6 TZ_NYC = pytz.timezone('America/New_York')
      7 TZ_UTC = pytz.timezone('UTC')
      8 
      9 logger = logging.getLogger('forerad')
     10 stream = logging.StreamHandler()
     11 fmt = logging.Formatter("%(asctime)s [%(levelname)s]: %(message)s")
     12 stream.setFormatter(fmt)
     13 logger.addHandler(stream)
     14 logger.setLevel(logging.INFO)
     15 
     16 def parse_month_str(month_str: str) -> tuple[int, int]:
     17     """
     18     Parses a string formatted YYYY-MM into a tuple of (year, month). Used for
     19     parsing CLI arguments to bin/ scripts
     20     """
     21     match = re.match(r"^([\d]{4})-([\d]{2})$", month_str)
     22     if match is None:
     23         raise Exception(f'Invalid month string: {month_str}')
     24 
     25     year, month = match.groups()
     26     return (int(year), int(month))
     27 
     28 def next_month(year, month) -> datetime.date:
     29     """
     30     Given a year and month, calculate the first day of the next month
     31     """
     32     date = datetime.date(year, month, 1)
     33     return (date + datetime.timedelta(days=32)).replace(day=1)