forerad

historical.py (7061B)

---

 import re
 import datetime
 import zipfile
 import io
 import boto3
 import pathlib
 import pandas as pd
 import numpy as np
 import botocore.client as bclient
 from botocore import UNSIGNED
 
 import forerad.utils as utils
 
 # Notice the `(i)?` - they spelled Citibike wrong in one of the archives... lol
 ARCHIVE_REGEX = re.compile("^([0-9]{4})([0-9]{2})-cit(i)?bike-tripdata((.zip$)|(.csv.zip$))")
 CACHE_DIR = pathlib.Path(__file__).parent.parent.parent / pathlib.Path('.forerad-cache')
 TRIP_BUCKET = 'tripdata'
 if not CACHE_DIR.exists():
     CACHE_DIR.mkdir()
     utils.logger.debug('Initializing .cache dir')
 
 
 class HistoricalTripArchive():
     object_key: str
     date: datetime.date
 
     def __init__(self, object_key, date):
         self.object_key = object_key
         self.date = date
 
     @classmethod
     def __get_s3_client(cls):
         config = bclient.Config(signature_version=UNSIGNED)
         return boto3.client('s3', config=config)
 
     @classmethod
     def from_s3(cls, obj: dict):
         """
         Converts an S3 object dictionary into a HistoricalTripArchive record, returning
         None if the object doesn't look like a monthly ride archive.
         """
         match = ARCHIVE_REGEX.match(obj['Key'])
         if match is None:
             utils.logger.error(f"Skipping object {obj['Key']}")
             return None
 
         groups = match.groups()
 
         return cls(
             object_key=obj['Key'],
             date=datetime.date(int(groups[0]), int(groups[1]), 1)
         )
 
     @classmethod
     def from_archive_path(cls, path: pathlib.Path):
         match = ARCHIVE_REGEX.match(path.name)
         if match is None:
             return None
 
         groups = match.groups()
 
         return cls(
             object_key=path.name,
             date=datetime.date(int(groups[0]), int(groups[1]), 1)
         )
 
     @classmethod
     def list_remote(cls):
         """
         Returns a list of HistoricalTripArchive objects that are available on
         S3 but may or may not be cached locally.
         """
         s3 = cls.__get_s3_client()
         resp = s3.list_objects_v2(Bucket=TRIP_BUCKET)
 
         archives = [HistoricalTripArchive.from_s3(o) for o in resp['Contents']]
         return [a for a in archives if a is not None]
 
     @classmethod
     def list_cached(cls):
         """
         Returns a list of HistoricalTripArchive objects for all archives that have already
         been downloaded
         """
         archives = [cls.from_archive_path(p) for p in CACHE_DIR.glob('*')]
         archives = [a for a in archives if a is not None]
         return sorted(archives, key=lambda a: a.date)
     
     def __fetch_cached_blob(self) -> io.BytesIO | None:
         archive_path = CACHE_DIR / self.object_key
         if not archive_path.exists():
             return None
 
         with open(archive_path, 'rb') as f:
             return io.BytesIO(f.read())
 
     def __store_blob(self, blob: io.BytesIO):
         archive_path = CACHE_DIR / self.object_key
         with open(archive_path, 'wb') as f:
             utils.logger.info(f"Storing {self.object_key} in cache")
             f.write(blob.getbuffer())
 
         blob.seek(0)
 
     def fetch_df(self, normalize=True) -> pd.DataFrame:
         """
         Fetches a DataFrame of the archive. If the archive doesn't exist on the
         filesystem it will attempt to download it from S3.
         """
         blob = self.__fetch_cached_blob()
         if blob is None:
             utils.logger.info(f"Fetching {self.csv_name} from S3")
             s3 = self.__get_s3_client()
             resp = s3.get_object(Bucket=TRIP_BUCKET, Key=self.object_key)
             blob = io.BytesIO(resp['Body'].read())
             self.__store_blob(blob)
 
         # Unzip the CSV into a dataframe
         with zipfile.ZipFile(blob, 'r') as zip_archive:
             file_list = zip_archive.namelist()
             csv_name = self.csv_name
 
             if csv_name not in file_list and len(file_list) != 1:
                 utils.logger.error(f"Could not extract {self.csv_name}:")
                 utils.logger.error(file_list)
                 raise Exception("Could not extract {self.csv_name}")
 
             if csv_name not in file_list:
                 csv_name = file_list[0]
 
             with zip_archive.open(csv_name, 'r') as csv:
                 df = pd.read_csv(csv)
 
         if not normalize:
             return df
 
         # Normalize the dataframe
         if 'starttime' in df.columns:
             # Transform v1 schema
             df = df.rename(columns={
                 'starttime': 'started_at',
                 'stoptime': 'ended_at',
                 'start station id': 'start_station_id',
                 'end station id': 'end_station_id',
                 'start station latitude': 'start_lat',
                 'start station longitude': 'start_lng',
                 'end station latitude': 'end_lat',
                 'end station longitude': 'end_lng',
             })
         elif 'Start Time' in df.columns:
             # This is a weird in-between state of v1 and v2 that starts in the
             # 2016-10 dataset
             df = df.rename(columns={
                 'Start Time': 'started_at',
                 'Stop Time': 'ended_at',
                 'Start Station ID': 'start_station_id',
                 'End Station ID': 'end_station_id',
                 'Start Station Latitude': 'start_lat',
                 'Start Station Longitude': 'start_lng',
                 'End Station Latitude': 'end_lat',
                 'End Station Longitude': 'end_lng',
             })
 
         if 'rideable_type' not in df.columns:
             df['rideable_type'] = None
 
         df = df[[
             'rideable_type', 'started_at', 'ended_at', 'start_station_id',
             'end_station_id', 'start_lat', 'start_lng', 'end_lat', 'end_lng',
         ]]
         df = df.replace({np.nan: None})
 
         # Timestamps in the archives are in NYC time but don't differentiate
         # between EST and EDT. I'm admittedly punting here by assigning
         # ambiguous times to NaT and then dropna'ing them below. At some point
         # this will likely need to be revisited.
         df['started_at'] = pd.to_datetime(df['started_at'])\
             .dt.tz_localize(utils.TZ_NYC, ambiguous='NaT')\
             .dt.tz_convert(utils.TZ_UTC)
         df['ended_at'] = pd.to_datetime(df['ended_at'])\
             .dt.tz_localize(utils.TZ_NYC, ambiguous='NaT')\
             .dt.tz_convert(utils.TZ_UTC)
 
         df.index.name = 'tempid'
 
         return df.dropna(subset=['started_at', 'ended_at'])
 
     @property
     def is_downloaded(self) -> bool:
         return (CACHE_DIR / self.object_key).exists()
 
     @property
     def month_str(self) -> str:
         return f'{self.date.year}-{self.date.month:02d}'
 
     @property
     def csv_name(self) -> str:
         if '.csv.zip' in self.object_key:
             return self.object_key.replace('.csv.zip', '.csv')
 
         return self.object_key.replace('.zip', '.csv')
 
     def __repr__(self):
         return f'<HistoricalTripArchive {self.date} {self.object_key} />'