From f796bacf3b07ccd85141b42eda34aa33de1c25f3 Mon Sep 17 00:00:00 2001
From: Becky Sweger <rsweger@umass.edu>
Date: Tue, 1 Oct 2024 13:41:06 -0400
Subject: [PATCH] More type checking fixes

---
 pyproject.toml                         | 6 +++++-
 requirements/requirements-dev.txt      | 1 +
 requirements/requirements.txt          | 1 +
 src/virus_clade_utils/util/sequence.py | 2 +-
 src/virus_clade_utils/util/session.py  | 4 ++--
 tests/unit/test_cladetime.py           | 4 ++--
 6 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 3a9383a..d02cda6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,8 @@ dependencies = [
     "rich",
     "rich-click",
     "structlog",
-    "us"
+    "urllib3",
+    "us",
 ]
 
 [project.optional-dependencies]
@@ -63,3 +64,6 @@ lint.extend-select = ["I"]
 [tools.setuptools]
 packages = ["virus_clade_utils"]
 
+[tool.mypy]
+ignore_missing_imports = true
+
diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt
index 28b4407..9596d67 100644
--- a/requirements/requirements-dev.txt
+++ b/requirements/requirements-dev.txt
@@ -145,6 +145,7 @@ tzdata==2024.1
     # via pandas
 urllib3==2.2.1
     # via
+    #   virus-clade-utils (pyproject.toml)
     #   botocore
     #   requests
     #   responses
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index de13f4a..ecd54a0 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -81,6 +81,7 @@ tzdata==2024.1
     # via pandas
 urllib3==2.2.1
     # via
+    #   virus-clade-utils (pyproject.toml)
     #   botocore
     #   requests
 us==3.2.0
diff --git a/src/virus_clade_utils/util/sequence.py b/src/virus_clade_utils/util/sequence.py
index dcbc6c1..1e3d14a 100644
--- a/src/virus_clade_utils/util/sequence.py
+++ b/src/virus_clade_utils/util/sequence.py
@@ -9,7 +9,7 @@
 
 import polars as pl
 import structlog
-import us  # type: ignore
+import us
 from requests import Session
 
 from virus_clade_utils.util.reference import _get_s3_object_url
diff --git a/src/virus_clade_utils/util/session.py b/src/virus_clade_utils/util/session.py
index aa4b1e8..a7a7b9d 100644
--- a/src/virus_clade_utils/util/session.py
+++ b/src/virus_clade_utils/util/session.py
@@ -3,12 +3,12 @@
 import requests
 import structlog
 from requests.adapters import HTTPAdapter
-from requests.packages.urllib3.util.retry import Retry  # type: ignore
+from urllib3.util import Retry
 
 logger = structlog.get_logger()
 
 
-def _get_session(session: requests.Session = None) -> requests.Session:
+def _get_session() -> requests.Session:
     """Return a requests session with retry logic."""
 
     headers = {
diff --git a/tests/unit/test_cladetime.py b/tests/unit/test_cladetime.py
index 507a6b9..5e7e4d7 100644
--- a/tests/unit/test_cladetime.py
+++ b/tests/unit/test_cladetime.py
@@ -5,8 +5,8 @@
 import dateutil.tz
 import pytest
 from freezegun import freeze_time
-from virus_clade_utils.cladetime import CladeTime  # type: ignore
-from virus_clade_utils.exceptions import CladeTimeInvalidDateError  # type: ignore
+from virus_clade_utils.cladetime import CladeTime
+from virus_clade_utils.exceptions import CladeTimeInvalidDateError
 from virus_clade_utils.util.config import Config