first draft of typing, typo-fixes, docstring updates

src/probeinterface/generator.py

@@ -10,15 +10,15 @@
 from .utils import combine_probes
 
 
-def generate_dummy_probe(elec_shapes="circle"):
+def generate_dummy_probe(elec_shapes:str ="circle"):
     """
     Generate a dummy probe with 3 columns and 32 contacts.
     Mainly used for testing and examples.
 
     Parameters
     ----------
-    elec_shapes : str, optional
-        Shape of the electrodes, by default 'circle'
+    elec_shapes : str, , by default 'circle'
+        Shape of the electrodes with possibilities of ('circle', 'square', 'rect')
 
     Returns
     -------
@@ -72,10 +72,13 @@ def generate_dummy_probe_group():
     return probegroup
 
 
-def generate_tetrode(r=10):
+def generate_tetrode(r:float=10):
     """
     Generate a tetrode Probe.
-
+    Parameters
+    ----------
+    r: float
+        The distance multiplier for the positions
     Returns
     -------
     probe : Probe
@@ -89,32 +92,35 @@ def generate_tetrode(r=10):
 
 
 def generate_multi_columns_probe(
-    num_columns=3,
-    num_contact_per_column=10,
-    xpitch=20,
-    ypitch=20,
+    num_columns:int=3,
+    num_contact_per_column:int =10,
+    xpitch:float=20,
+    ypitch:float=20,
     y_shift_per_column=None,
-    contact_shapes="circle",
-    contact_shape_params={"radius": 6},
+    contact_shapes:str="circle",
+    contact_shape_params:dict={"radius": 6},
 ):
     """Generate a Probe with several columns.
 
     Parameters
     ----------
-    num_columns : int, optional
-        Number of columns, by default 3
-    num_contact_per_column : int, optional
-        Number of contacts per column, by default 10
-    xpitch : float, optional
-        Pitch in x direction, by default 20
-    ypitch : float, optional
-        Pitch in y direction, by default 20
+    num_columns : int, by default 3
+        Number of columns
+    num_contact_per_column : int, by default 10
+        Number of contacts per column
+    xpitch : float, by default 20
+        Pitch in x direction
+    ypitch : float, by default 20
+        Pitch in y direction
     y_shift_per_column : array-like, optional
         Shift in y direction per column. It needs to have the same length as num_columns, by default None
-    contact_shapes : str, optional
-        Shape of the contacts ('circle', 'rect', 'square'), by default 'circle'
-    contact_shape_params : dict, optional
-        Parameters for the shape, by default {'radius': 6}
+    contact_shapes : str, by default 'circle'
+        Shape of the contacts ('circle', 'rect', 'square')
+    contact_shape_params : dict, default {'radius': 6}
+        Parameters for the shape.
+        For circle: {"radius": float}
+        For square: {"width": float}
+        For rectangle: {"width": float, "height": float}
 
     Returns
     -------
@@ -144,19 +150,22 @@ def generate_multi_columns_probe(
     return probe
 
 
-def generate_linear_probe(num_elec=16, ypitch=20, contact_shapes="circle", contact_shape_params={"radius": 6}):
+def generate_linear_probe(num_elec: int =16, ypitch:float=20, contact_shapes:str="circle", contact_shape_params:dict={"radius": 6}):
     """Generate a one-column linear probe.
 
     Parameters
     ----------
-    num_elec : int, optional
+    num_elec : int
         Number of electrodes, by default 16
-    ypitch : float, optional
+    ypitch : float
         Pitch in y direction, by default 20
-    contact_shapes : str, optional
-        Shape of the contacts ('circle', 'rect', 'square'), by default 'circle'
-    contact_shape_params : dict, optional
-        Parameters for the shape, by default {'radius': 6}
+    contact_shapes : str, default 'circle'
+        Shape of the contacts ('circle', 'rect', 'square')
+    contact_shape_params : dict, default {'radius': 6}
+        Parameters for the shape.
+        For circle: {"radius": float}
+        For square: {"width": float}
+        For rectangle: {"width": float, "height": float}
 
     Returns
     -------
@@ -175,15 +184,15 @@ def generate_linear_probe(num_elec=16, ypitch=20, contact_shapes="circle", conta
     return probe
 
 
-def generate_multi_shank(num_shank=2, shank_pitch=[150, 0], **kargs):
+def generate_multi_shank(num_shank:int=2, shank_pitch:list=[150, 0], **kargs):
     """Generate a multi-shank probe.
     Internally, calls generate_multi_columns_probe and combine_probes.
 
     Parameters
     ----------
-    num_shank : int, optional
-        Number of shanks, by default 2
-    shank_pitch : list, optional
+    num_shank : int, default 2
+        Number of shanks
+    shank_pitch : list, default [150,0]
         Distance between shanks, by default [150, 0]
 
     Returns

src/probeinterface/io.py

@@ -32,7 +32,7 @@ def _probeinterface_format_check_version(d):
     pass
 
 
-def read_probeinterface(file):
+def read_probeinterface(file:Union[str, Path])-> ProbeGroup:
     """
     Read probeinterface JSON-based format.
 
@@ -80,7 +80,7 @@ def write_probeinterface(file: Union[str, Path], probe_or_probegroup: Union[Prob
     elif isinstance(probe_or_probegroup, ProbeGroup):
         probegroup = probe_or_probegroup
     else:
-        raise ValueError("write_probeinterface : need probe or probegroup")
+        raise ValueError("write_probeinterface : needs a probe or probegroup")
 
     file = Path(file)
 
@@ -175,23 +175,23 @@ def read_BIDS_probe(folder: Union[str, Path], prefix: Optional[str] = None) -> P
         if "contact_shapes" not in df_probe:
             df_probe["contact_shapes"] = "circle"
             df_probe["radius"] = 1
-            print(f"There is no contact shape provided for probe {probe_id}, a " f"dummy circle with 1um is created")
+            print(f"There is no contact shape provided for probe {probe_id}, a " f"dummy circle with 1um radius will be used.")
 
         if "x" not in df_probe:
             df_probe["x"] = np.arange(len(df_probe.index), dtype=float)
             print(
-                f"There is no x coordinate provided for probe {probe_id}, a " f"dummy linear x coordinate is created."
+                f"There is no x coordinate provided for probe {probe_id}, a " f"dummy linear x coordinate will be used."
             )
 
         if "y" not in df_probe:
             df_probe["y"] = 0.0
             print(
-                f"There is no y coordinate provided for probe {probe_id}, a " f"dummy constant y coordinate is created."
+                f"There is no y coordinate provided for probe {probe_id}, a " f"dummy constant y coordinate will be used."
             )
 
         if "si_units" not in df_probe:
             df_probe["si_units"] = "um"
-            print(f"There is no SI units provided for probe {probe_id}, a " f"dummy SI unit (um) is created.")
+            print(f"There is no SI unit provided for probe {probe_id}, a " f"dummy SI unit (um) will be used")
 
         # create probe object and register with probegroup
         probe = Probe.from_dataframe(df=df_probe)
@@ -290,7 +290,7 @@ def read_BIDS_probe(folder: Union[str, Path], prefix: Optional[str] = None) -> P
     return probegroup
 
 
-def write_BIDS_probe(folder: Union[str, Path], probe_or_probegroup: Union[Probe, ProbeGroup], prefix=""):
+def write_BIDS_probe(folder: Union[str, Path], probe_or_probegroup: Union[Probe, ProbeGroup], prefix:str=""):
     """
     Write to probe and contact formats as proposed
     for ephy BIDS extension (tsv & json based).
@@ -577,7 +577,7 @@ def read_3brain(file: Union[str, Path], mea_pitch: float = 42, electrode_width:
     return probe
 
 
-def write_prb(file, probegroup, total_nb_channels=None, radius=None, group_mode="by_probe"):
+def write_prb(file:str, probegroup:ProbeGroup, total_nb_channels:Optional[int]=None, radius:Optional[float]=None, group_mode:str="by_probe"):
     """
     Write ProbeGroup into a prb file.
 
@@ -596,6 +596,19 @@ def write_prb(file, probegroup, total_nb_channels=None, radius=None, group_mode=
       * "radius" is needed by spyking-circus
       * "graph" is not handled
 
+    Parameters
+    ----------
+    file: str
+        The name of the file to be written
+    probegroup: ProbeGroup
+        The Probegroup to be used for writing
+    total_nb_channels: Optional[int], default None
+        ***to do
+    radius: Optional[float], default None
+        *** to do
+    group_mode: str
+        One of "by_probe" or "by_shank
+
     """
     assert group_mode in ("by_probe", "by_shank")
 
@@ -643,7 +656,7 @@ def write_prb(file, probegroup, total_nb_channels=None, radius=None, group_mode=
         f.write("}\n")
 
 
-def read_csv(file):
+def read_csv(file:Union[str, Path]):
     """
     Return a 2 or 3 columns csv file with contact positions
     """

src/probeinterface/library.py

@@ -11,6 +11,7 @@
 import os
 from pathlib import Path
 from urllib.request import urlopen
+from typing import Optional
 
 from .io import read_probeinterface
 
@@ -24,7 +25,7 @@
 cache_folder = Path(os.path.expanduser("~")) / ".config" / "probeinterface" / "library"
 
 
-def download_probeinterface_file(manufacturer, probe_name):
+def download_probeinterface_file(manufacturer:str, probe_name:str):
     """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.
@@ -44,14 +45,14 @@ def download_probeinterface_file(manufacturer, probe_name):
         f.write(dist.read())
 
 
-def get_from_cache(manufacturer, probe_name):
+def get_from_cache(manufacturer:str, probe_name:str)-> Optional["Probe"]:
     """
     Get Probe from local cache
 
     Parameters
     ----------
     manufacturer : str
-        The probe manufacturer (e.g. 'cambridgeneurotech')
+        The probe manufacturer (e.g. 'cambridgeneurotech', 'neuronexus')
     probe_name : str
         The probe name
 
@@ -71,14 +72,14 @@ def get_from_cache(manufacturer, probe_name):
         return probe
 
 
-def get_probe(manufacturer, probe_name):
+def get_probe(manufacturer:str, probe_name:str)-> "Probe":
     """
     Get probe from ProbeInterface library
 
     Parameters
     ----------
     manufacturer : str
-        The probe manufacturer (e.g. 'cambridgeneurotech')
+        The probe manufacturer (e.g. 'cambridgeneurotech', 'neuronexus')
     probe_name : str
         The probe name