Refactor: module for I/O

air_traffic/io.py

@@ -0,0 +1,33 @@
+import pandas as pd
+import os
+
+
+def read_trajectories_range(datadir, start, end):
+
+    start_date = pd.to_datetime(start)
+    end_date = pd.to_datetime(end)
+    dt = pd.Timedelta(1, "D")
+
+    date = start_date
+    while date <= end_date:
+
+        dstr = date.strftime("%Y-%m-%d")
+        mstr = date.strftime(r"%Y-%m")
+
+        date += dt
+
+        # Skip over months with no data
+        if not os.path.exists(f"{datadir}/{mstr}"):
+            continue
+
+        yield read_trajectories(datadir, dstr)
+
+
+def read_trajectories(datadir, date):
+
+    path = os.path.join(datadir, date[:7], date)
+
+    for subdir, dirs, files in os.walk(path):
+        for file in files:
+            fname = os.path.join(subdir, file)
+            yield pd.read_csv(fname, header=0, index_col=0)

tests/test_io.py

@@ -0,0 +1,25 @@
+from air_traffic.io import *
+
+
+def test_read_trajectories():
+
+    dfs = read_trajectories("../data/cleaned", "2017-01-01")
+    count = 0
+    for df in dfs:
+        count += 1
+    print(count)
+
+
+def test_read_trajectories_range():
+
+    datasets = read_trajectories_range("../data/cleaned", "2017-01-01", "2017-01-02")
+
+    for dfs in datasets:
+        count = 0
+        for df in dfs:
+            count += 1
+        print(count)
+
+
+# test_read_trajectories()
+test_read_trajectories_range()