Merge pull request #5 from nolanbconaway/add-tz-option

add tz option for extract stop dict.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.masonry.api"
 
 [tool.poetry]
 name = "underground"
-version = "0.2.1"
+version = "0.2.1.1"
 description = ""
 authors = ["Nolan Conaway <nolanbconaway@gmail.com>"]
 packages = [{ include = "underground" }]

underground/cli/stops.py

@@ -1,7 +1,6 @@
 """Get upcoming stops along a train route."""
 
 import click
-import pytz
 
 from underground import dateutils, metadata
 from underground.models import SubwayFeed
@@ -45,15 +44,15 @@ def main(route, fmt, retries, api_key, timezone):
         SubwayFeed.get(
             api_key=api_key, feed_id=metadata.ROUTE_FEED_MAP.get(route), retries=retries
         )
-        .extract_stop_dict()
+        .extract_stop_dict(timezone=timezone)
         .get(route, dict())
     )
 
     # figure out how to format it
     if fmt == "epoch":
         format_fun = dateutils.datetime_to_epoch
     else:
-        format_fun = lambda x: x.astimezone(pytz.timezone(timezone)).strftime(fmt)
+        format_fun = lambda x: x.strftime(fmt)
 
     # echo the result
     for stop_id, departures in stops.items():

underground/models.py

@@ -143,9 +143,14 @@ def get(feed_id: int, retries: int = 100, api_key: str = None) -> "SubwayFeed":
             )
         )
 
-    def extract_stop_dict(self) -> dict:
+    def extract_stop_dict(self, timezone: str = dateutils.DEFAULT_TIMEZONE) -> dict:
         """Get the departure times for all stops in the feed.
 
+        Parameters
+        ----------
+        timezone : str
+            Name of the timezone to return within. Default to NYC time.
+
         Returns
         -------
         dict
@@ -164,7 +169,11 @@ def extract_stop_dict(self) -> dict:
         )
         # create (route, stop, time) tuples from each trip
         stops_flat = (
-            (trip.trip.route_id, stop.stop_id, stop.departure.time)
+            (
+                trip.trip.route_id,
+                stop.stop_id,
+                stop.departure.time.astimezone(pytz.timezone(timezone)),
+            )
             for trip in trip_updates_with_stops
             for stop in trip.stop_time_update
         )