Fix annotations

Author: alejoe91

src/probeinterface/generator.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from typing import Optional
+from typing import Literal
 
 from .probe import Probe
 from .probegroup import ProbeGroup
@@ -14,7 +14,7 @@
 _default_shape_to_params = {"circle": "radius", "square": "width", "rect": "height"}
 
 
-def generate_dummy_probe(elec_shapes: "circle" | "square" | "rect" = "circle") -> Probe:
+def generate_dummy_probe(elec_shapes: Literal["circle", "square", "rect"] = "circle") -> Probe:
     """
     Generate a dummy probe with 3 columns and 32 contacts.
     Mainly used for testing and examples.
@@ -102,8 +102,8 @@ def generate_multi_columns_probe(
     num_contact_per_column: int | list[int] = 10,
     xpitch: float = 20,
     ypitch: float = 20,
-    y_shift_per_column: Optional[np.array | list] = None,
-    contact_shapes: "circle" | "rect" | "square" = "circle",
+    y_shift_per_column: np.ndarray | list | None = None,
+    contact_shapes: Literal["circle", "rect", "square"] = "circle",
     contact_shape_params: dict = {"radius": 6},
 ) -> Probe:
     """Generate a Probe with several columns.
@@ -118,7 +118,7 @@ def generate_multi_columns_probe(
         Pitch in x direction
     ypitch : float, default: 20
         Pitch in y direction
-    y_shift_per_column : Optional[array-like], default: None
+    y_shift_per_column : array-like | None, default: None
         Shift in y direction per column. It needs to have the same length as num_columns, by default None
     contact_shapes : "circle" | "rect" | "square", default: "circle"
         Shape of the contacts
@@ -166,7 +166,7 @@ def generate_multi_columns_probe(
 def generate_linear_probe(
     num_elec: int = 16,
     ypitch: float = 20,
-    contact_shapes: "circle" | "rect" | "square" = "circle",
+    contact_shapes: Literal["circle", "rect", "square"] = "circle",
     contact_shape_params: dict = {"radius": 6},
 ) -> Probe:
     """Generate a one-column linear probe.

src/probeinterface/io.py

@@ -10,7 +10,6 @@
 """
 
 from pathlib import Path
-from typing import Union, Optional
 import re
 import warnings
 import json
@@ -107,7 +106,7 @@ def write_probeinterface(file: str | Path, probe_or_probegroup: Probe | ProbeGro
 tsv_label_map_to_probeinterface = {v: k for k, v in tsv_label_map_to_BIDS.items()}
 
 
-def read_BIDS_probe(folder: str | Path, prefix: Optional[str] = None) -> ProbeGroup:
+def read_BIDS_probe(folder: str | Path, prefix: str | None = None) -> ProbeGroup:
     """
     Read to BIDS probe format.
 
@@ -635,8 +634,8 @@ def read_3brain(file: str | Path, mea_pitch: float = None, electrode_width: floa
 def write_prb(
     file: str,
     probegroup: ProbeGroup,
-    total_nb_channels: Optional[int] = None,
-    radius: Optional[float] = None,
+    total_nb_channels: int | None = None,
+    radius: float | None = None,
     group_mode: str = "by_probe",
 ):
     """
@@ -663,9 +662,9 @@ def write_prb(
         The name of the file to be written
     probegroup: ProbeGroup
         The Probegroup to be used for writing
-    total_nb_channels: Optional[int], default None
+    total_nb_channels: int | None, default None
         ***to do
-    radius: Optional[float], default None
+    radius: float | None, default None
         *** to do
     group_mode: str
         One of "by_probe" or "by_shank

src/probeinterface/library.py

@@ -14,9 +14,9 @@
 from pathlib import Path
 from urllib.request import urlopen
 import requests
-from typing import Optional
 
 from .io import read_probeinterface
+from .probe import Probe
 
 # OLD URL on gin
 # public_url = "https://web.gin.g-node.org/spikeinterface/probeinterface_library/raw/master/"
@@ -36,7 +36,7 @@ def get_cache_folder() -> Path:
     return Path(os.path.expanduser("~")) / ".config" / "probeinterface" / "library"
 
 
-def download_probeinterface_file(manufacturer: str, probe_name: str, tag: Optional[str] = None) -> None:
+def download_probeinterface_file(manufacturer: str, probe_name: str, tag: str | None = None) -> None:
     """Download the probeinterface file to the cache directory.
     Note that the file is itself a ProbeGroup but on the repo each file
     represents one probe.
@@ -64,7 +64,7 @@ def download_probeinterface_file(manufacturer: str, probe_name: str, tag: Option
         f.write(rem.read())
 
 
-def get_from_cache(manufacturer: str, probe_name: str, tag: Optional[str] = None) -> Optional["Probe"]:
+def get_from_cache(manufacturer: str, probe_name: str, tag: str | None = None) -> Probe | None:
     """
     Get Probe from local cache
 
@@ -104,8 +104,8 @@ def get_from_cache(manufacturer: str, probe_name: str, tag: Optional[str] = None
 def get_probe(
     manufacturer: str,
     probe_name: str,
-    name: Optional[str] = None,
-    tag: Optional[str] = None,
+    name: str | None = None,
+    tag: str | None = None,
     force_download: bool = False,
 ) -> "Probe":
     """

src/probeinterface/neuropixels_tools.py

@@ -8,7 +8,6 @@
 """
 
 from pathlib import Path
-from typing import Union, Optional
 import warnings
 from packaging.version import parse
 import json
@@ -113,7 +112,7 @@ def get_probe_length(probe_part_number: str) -> int:
     return 10_000
 
 
-def make_mux_table_array(mux_information) -> np.array:
+def make_mux_table_array(mux_information) -> np.ndarray:
     """
     Function to parse the mux_table from ProbeTable.
 
@@ -209,7 +208,7 @@ def get_probe_contour_vertices(shank_width, tip_length, probe_length) -> list:
     return polygon_vertices
 
 
-def read_imro(file_path: Union[str, Path]) -> Probe:
+def read_imro(file_path: str | Path) -> Probe:
     """
     Read probe position from the imro file used in input of SpikeGlx and Open-Ephys for neuropixels probes.
 
@@ -570,7 +569,7 @@ def _annotate_probe_with_adc_sampling_info(probe: Probe, adc_sampling_table: str
     _annotate_contacts_from_mux_table(probe, adc_groups_array)
 
 
-def _read_imro_string(imro_str: str, imDatPrb_pn: Optional[str] = None) -> Probe:
+def _read_imro_string(imro_str: str, imDatPrb_pn: str | None = None) -> Probe:
     """
     Parse the IMRO table when presented as a string and create a Probe object.
 
@@ -1026,7 +1025,7 @@ def parse_spikeglx_snsGeomMap(meta):
 #     parse_spikeglx_snsGeomMap(meta)
 
 
-def get_saved_channel_indices_from_spikeglx_meta(meta_file: str | Path) -> np.array:
+def get_saved_channel_indices_from_spikeglx_meta(meta_file: str | Path) -> np.ndarray:
     """
     Utils function to get the saved channels.
 
@@ -1066,7 +1065,7 @@ def _parse_openephys_settings(
     settings_file: str | Path,
     fix_x_position_for_oe_5: bool = True,
     raise_error: bool = True,
-) -> Optional[list[dict]]:
+) -> list[dict] | None:
     """
     Parse an Open Ephys settings.xml and extract per-probe metadata.
 
@@ -1395,11 +1394,11 @@ def _parse_openephys_settings(
 
 def _select_openephys_probe_info(
     probes_info: list[dict],
-    stream_name: Optional[str] = None,
-    probe_name: Optional[str] = None,
-    serial_number: Optional[str] = None,
+    stream_name: str | None = None,
+    probe_name: str | None = None,
+    serial_number: str | None = None,
     raise_error: bool = True,
-) -> Optional[dict]:
+) -> dict | None:
     """
     Select one probe's info dict from the list returned by `_parse_openephys_settings`.
 
@@ -1586,9 +1585,9 @@ def _annotate_openephys_probe(probe: Probe, probe_info: dict) -> None:
 
 def read_openephys(
     settings_file: str | Path,
-    stream_name: Optional[str] = None,
-    probe_name: Optional[str] = None,
-    serial_number: Optional[str] = None,
+    stream_name: str | None = None,
+    probe_name: str | None = None,
+    serial_number: str | None = None,
     fix_x_position_for_oe_5: bool = True,
     raise_error: bool = True,
 ) -> Probe:
@@ -1675,9 +1674,7 @@ def read_openephys(
     return probe
 
 
-def get_saved_channel_indices_from_openephys_settings(
-    settings_file: str | Path, stream_name: str
-) -> Optional[np.array]:
+def get_saved_channel_indices_from_openephys_settings(settings_file: str | Path, stream_name: str) -> np.ndarray | None:
     """
     Returns an array with the subset of saved channels indices (if used)