Apply ruff fixes.

Also needed a mypy ignore on the extract monstrosity.

4 files changed, 35 insertions, 30 deletions

diff --git a/notebooks/cc/reference_annotations/hyperparameter_search.py b/notebooks/cc/reference_annotations/hyperparameter_search.py
index a7286aa0..fc42d148 100644
--- a/notebooks/cc/reference_annotations/hyperparameter_search.py
+++ b/notebooks/cc/reference_annotations/hyperparameter_search.py
@@ -14,11 +14,11 @@ import optuna
 import pandas as pd
 import torch
 from rapidfuzz import fuzz
+from sec_certs.model.references.annotator_trainer import ReferenceAnnotatorTrainer
+from sec_certs.model.references.segment_extractor import ReferenceSegmentExtractor
 from sklearn.metrics import f1_score
 
 from sec_certs.dataset import CCDataset
-from sec_certs.model.references.annotator_trainer import ReferenceAnnotatorTrainer
-from sec_certs.model.references.segment_extractor import ReferenceSegmentExtractor
 from sec_certs.utils.helpers import compute_heuristics_version
 from sec_certs.utils.nlp import prec_recall_metric
 
diff --git a/notebooks/fixed_sankey_plot.py b/notebooks/fixed_sankey_plot.py
index b8d062f9..8f609227 100644
--- a/notebooks/fixed_sankey_plot.py
+++ b/notebooks/fixed_sankey_plot.py
@@ -9,7 +9,7 @@ This code should fix the problems and should be used to produce figures in the r
 import logging
 import warnings
 from collections import defaultdict
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -57,18 +57,18 @@ def check_data_matches_labels(labels: Union[list[str], set[str]], data: Series,
 def sankey(
     left: Union[list, ndarray, Series],
     right: Union[ndarray, Series],
-    leftWeight: Optional[ndarray] = None,
-    rightWeight: Optional[ndarray] = None,
-    colorDict: Optional[dict[str, str]] = None,
-    leftLabels: Optional[list[str]] = None,
-    rightLabels: Optional[list[str]] = None,
+    leftWeight: ndarray | None = None,
+    rightWeight: ndarray | None = None,
+    colorDict: dict[str, str] | None = None,
+    leftLabels: list[str] | None = None,
+    rightLabels: list[str] | None = None,
     aspect: int = 4,
     rightColor: bool = False,
     fontsize: int = 14,
-    figureName: Optional[str] = None,
+    figureName: str | None = None,
     closePlot: bool = False,
-    figSize: Optional[tuple[int, int]] = None,
-    ax: Optional[Any] = None,
+    figSize: tuple[int, int] | None = None,
+    ax: Any | None = None,
 ) -> Any:
     """
     Make Sankey Diagram showing flow from left-->right
@@ -151,7 +151,7 @@ def sankey(
     return ax
 
 
-def save_image(figureName: Optional[str]) -> None:
+def save_image(figureName: str | None) -> None:
     if figureName is not None:
         file_name = f"{figureName}.png"
         plt.savefig(file_name, bbox_inches="tight", dpi=150)
@@ -173,15 +173,15 @@ def identify_labels(dataFrame: DataFrame, leftLabels: list[str], rightLabels: li
 
 
 def init_values(
-    ax: Optional[Any],
+    ax: Any | None,
     closePlot: bool,
-    figSize: Optional[tuple[int, int]],
-    figureName: Optional[str],
+    figSize: tuple[int, int] | None,
+    figureName: str | None,
     left: Union[list, ndarray, Series],
-    leftLabels: Optional[list[str]],
-    leftWeight: Optional[ndarray],
-    rightLabels: Optional[list[str]],
-    rightWeight: Optional[ndarray],
+    leftLabels: list[str] | None,
+    leftWeight: ndarray | None,
+    rightLabels: list[str] | None,
+    rightWeight: ndarray | None,
 ) -> tuple[Any, list[str], ndarray, list[str], ndarray]:
     deprecation_warnings(closePlot, figSize, figureName)
     if ax is None:
@@ -202,7 +202,7 @@ def init_values(
     return ax, leftLabels, leftWeight, rightLabels, rightWeight
 
 
-def deprecation_warnings(closePlot: bool, figSize: Optional[tuple[int, int]], figureName: Optional[str]) -> None:
+def deprecation_warnings(closePlot: bool, figSize: tuple[int, int] | None, figureName: str | None) -> None:
     warn = []
     if figureName is not None:
         msg = "use of figureName in sankey() is deprecated"
@@ -286,7 +286,7 @@ def draw_vertical_bars(
 
 
 def create_colors(
-    allLabels: ndarray, colorDict: Optional[dict[str, str]]
+    allLabels: ndarray, colorDict: dict[str, str] | None
 ) -> Union[dict[str, tuple[float, float, float]], dict[str, str]]:
     # If no colorDict given, make one
     if colorDict is None:
diff --git a/src/sec_certs/configuration.py b/src/sec_certs/configuration.py
index 17db9318..c06f872a 100644
--- a/src/sec_certs/configuration.py
+++ b/src/sec_certs/configuration.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 from pathlib import Path
-from typing import Literal, Optional
+from typing import Literal
 
 import yaml
 from pydantic import AnyHttpUrl, Field
@@ -121,7 +121,7 @@ class Configuration(BaseSettings):
         True,
         description="During keyword search, first page usually contains addresses - ignore it.",
     )
-    cc_reference_annotator_dir: Optional[Path] = Field(  # noqa: UP007
+    cc_reference_annotator_dir: Path | None = Field(  # noqa: UP007
         None,
         description="Path to directory with serialized reference annotator model. If set to `null`, tool will search default directory for the given dataset.",
     )
@@ -141,7 +141,7 @@ class Configuration(BaseSettings):
         True,
         description="If true, progress bars will be printed to stdout during computation.",
     )
-    nvd_api_key: Optional[str] = Field(None, description="NVD API key for access to CVEs and CPEs.")  # noqa: UP007
+    nvd_api_key: str | None = Field(None, description="NVD API key for access to CVEs and CPEs.")  # noqa: UP007
     preferred_source_remote_datasets: Literal["sec-certs", "origin"] = Field(
         "sec-certs",
         description="If set to `sec-certs`, will fetch remote datasets from sec-certs.org."
diff --git a/src/sec_certs/utils/extract.py b/src/sec_certs/utils/extract.py
index 35361e94..f72fd837 100644
--- a/src/sec_certs/utils/extract.py
+++ b/src/sec_certs/utils/extract.py
@@ -6,7 +6,7 @@ import re
 from collections import Counter
 from enum import Enum
 from pathlib import Path
-from typing import Any
+from typing import Any, no_type_check
 
 import numpy as np
 
@@ -17,7 +17,8 @@ from sec_certs.constants import FILE_ERRORS_STRATEGY, LINE_SEPARATOR, MAX_ALLOWE
 logger = logging.getLogger(__name__)
 
 
-def search_only_headers_anssi(filepath: Path):  # noqa: C901
+@no_type_check
+def search_only_headers_anssi(filepath: Path):  # type: ignore # noqa: C901
     # TODO: Please, refactor me. I reallyyyyyyyyyyyyy need it!!!!!!
     class HEADER_TYPE(Enum):
         HEADER_FULL = 1
@@ -272,7 +273,8 @@ def search_only_headers_anssi(filepath: Path):  # noqa: C901
     return items_found
 
 
-def search_only_headers_bsi(filepath: Path):  # noqa: C901
+@no_type_check
+def search_only_headers_bsi(filepath: Path):  # type: ignore # noqa: C901
     # TODO: Please, refactor me. I reallyyyyyyyyyyyyy need it!!!!!!
     LINE_SEPARATOR_STRICT = " "
     NUM_LINES_TO_INVESTIGATE = 15
@@ -371,7 +373,8 @@ def search_only_headers_bsi(filepath: Path):  # noqa: C901
     return items_found
 
 
-def search_only_headers_nscib(filepath: Path):  # noqa: C901
+@no_type_check
+def search_only_headers_nscib(filepath: Path):  # type: ignore # noqa: C901
     # TODO: Please, refactor me. I reallyyyyyyyyyyyyy need it!!!!!!
     LINE_SEPARATOR_STRICT = " "
     NUM_LINES_TO_INVESTIGATE = 60
@@ -451,7 +454,8 @@ def search_only_headers_nscib(filepath: Path):  # noqa: C901
     return items_found
 
 
-def search_only_headers_niap(filepath: Path):
+@no_type_check
+def search_only_headers_niap(filepath: Path):  # type: ignore # noqa: C901
     # TODO: Please, refactor me. I reallyyyyyyyyyyyyy need it!!!!!!
     LINE_SEPARATOR_STRICT = " "
     NUM_LINES_TO_INVESTIGATE = 15
@@ -502,7 +506,8 @@ def search_only_headers_niap(filepath: Path):
     return items_found
 
 
-def search_only_headers_canada(filepath: Path):  # noqa: C901
+@no_type_check
+def search_only_headers_canada(filepath: Path):  # type: ignore # noqa: C901
     # TODO: Please, refactor me. I reallyyyyyyyyyyyyy need it!!!!!!
     LINE_SEPARATOR_STRICT = " "
     NUM_LINES_TO_INVESTIGATE = 20