fix typing

xoolive · Nov 2, 2024 · 4fbf326 · 4fbf326
1 parent 2d942f4
commit 4fbf326
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 15 deletions.
diff --git a/src/traffic/plugins/bluesky.py b/src/traffic/plugins/bluesky.py
@@ -27,7 +27,7 @@ def to_bluesky(
     """Generates a Bluesky scenario file."""
 
     if minimum_time is not None:
-        traffic = traffic.after(minimum_time)
+        traffic = traffic.after(minimum_time)  # type: ignore
 
     if isinstance(filename, str):
         filename = Path(filename)

diff --git a/tests/test_clustering.py b/tests/test_clustering.py
@@ -21,13 +21,11 @@ def predict(self, X: Any) -> None:
 
 def test_clustering() -> None:
     switzerland = cast(Traffic, get_sample(collections, "switzerland"))
+    between = switzerland.between("2018-08-01 12:00", "2018-08-01 14:00")
 
-    smaller = cast(
-        Traffic,
-        switzerland.between("2018-08-01 12:00", "2018-08-01 14:00")
-        .assign_id()
-        .eval(max_workers=4),
-    )
+    assert between is not None
+
+    smaller = cast(Traffic, between.assign_id().eval(max_workers=4))
 
     t_clustering = smaller.clustering(
         nb_samples=15,

diff --git a/tests/test_cpa.py b/tests/test_cpa.py
@@ -9,11 +9,9 @@
 def test_cpa() -> None:
     switzerland = cast(Traffic, get_sample(collections, "switzerland"))
 
-    smaller = (
-        switzerland.between("2018-08-01 12:00", "2018-08-01 14:00")
-        .assign_id()
-        .eval()
-    )
+    between = switzerland.between("2018-08-01 12:00", "2018-08-01 14:00")
+    assert between is not None
+    smaller = between.assign_id().eval()
     assert smaller is not None
 
     cpa = smaller.closest_point_of_approach(

diff --git a/tests/test_generation.py b/tests/test_generation.py
@@ -36,9 +36,10 @@ def test_generation() -> None:
     def compute_timedelta(df: pd.DataFrame) -> pd.Series:
         return (df.timestamp - df.timestamp.min()).dt.total_seconds()
 
+    between = switzerland.between("2018-08-01 12:00", "2018-08-01 14:00")
+    assert between is not None
     smaller = (
-        switzerland.between("2018-08-01 12:00", "2018-08-01 14:00")
-        .assign_id()
+        between.assign_id()
         .resample(10)
         .compute_xy(projection=EuroPP())
         .assign(timedelta=compute_timedelta)

diff --git a/tests/test_traffic.py b/tests/test_traffic.py
@@ -131,7 +131,9 @@ def test_chaining_with_lambda() -> None:
 
 def test_none() -> None:
     assert switzerland.query("altitude > 60000") is None
-    assert switzerland.after("2018-08-01").before("2018-08-01") is None
+    after = switzerland.after("2018-08-01")
+    assert after is not None
+    assert after.before("2018-08-01") is None
     assert switzerland.iterate_lazy().query("altitude > 60000").eval() is None