Make docstrings-style compliant and add assert messaging

src/probeinterface/generator.py

@@ -12,15 +12,18 @@
 from .utils import combine_probes
 
 
-def generate_dummy_probe(elec_shapes: str = "circle") -> Probe:
+_default_shape_to_params = {"circle": "radius", "square": "width", "rect": "height"}
+
+
+def generate_dummy_probe(elec_shapes: "circle" | "square" | "rect" = "circle") -> Probe:
     """
     Generate a dummy probe with 3 columns and 32 contacts.
     Mainly used for testing and examples.
 
     Parameters
     ----------
-    elec_shapes : str, , by default 'circle'
-        Shape of the electrodes with possibilities of ('circle', 'square', 'rect')
+    elec_shapes : "circle" | "square" | "rect", default: 'circle'
+        Shape of the electrodes
 
     Returns
     -------
@@ -74,13 +77,15 @@ def generate_dummy_probe_group() -> ProbeGroup:
     return probegroup
 
 
-def generate_tetrode(r: float = 10) -> Probe:
+def generate_tetrode(r: float = 10.0) -> Probe:
     """
     Generate a tetrode Probe.
+
     Parameters
     ----------
-    r: float
+    r: float, default: 10
         The distance multiplier for the positions
+
     Returns
     -------
     probe : Probe
@@ -99,26 +104,26 @@ def generate_multi_columns_probe(
     xpitch: float = 20,
     ypitch: float = 20,
     y_shift_per_column: Optional[np.array | list] = None,
-    contact_shapes: str = "circle",
+    contact_shapes: "circle" | "rect" | "square" = "circle",
     contact_shape_params: dict = {"radius": 6},
 ) -> Probe:
     """Generate a Probe with several columns.
 
     Parameters
     ----------
-    num_columns : int, by default 3
+    num_columns : int, default: 3
         Number of columns
-    num_contact_per_column : int, by default 10
+    num_contact_per_column : int, default: 10
         Number of contacts per column
-    xpitch : float, by default 20
+    xpitch : float, default: 20
         Pitch in x direction
-    ypitch : float, by default 20
+    ypitch : float, default: 20
         Pitch in y direction
-    y_shift_per_column : array-like, optional
+    y_shift_per_column : Optional[array-like], default: None
         Shift in y direction per column. It needs to have the same length as num_columns, by default None
-    contact_shapes : str, by default 'circle'
-        Shape of the contacts ('circle', 'rect', 'square')
-    contact_shape_params : dict, default {'radius': 6}
+    contact_shapes : "circle" | "rect" | "square", default: "circle"
+        Shape of the contacts
+    contact_shape_params : dict, default: {'radius': 6}
         Parameters for the shape.
         For circle: {"radius": float}
         For square: {"width": float}
@@ -130,13 +135,19 @@ def generate_multi_columns_probe(
         The generated probe
     """
 
+    assert (
+        _default_shape_to_params[contact_shapes] in contact_shape_params.keys()
+    ), "contact_shapes and contact_shape_params must be coordinated see docstring"
+
     if isinstance(num_contact_per_column, int):
         num_contact_per_column = [num_contact_per_column] * num_columns
 
     if y_shift_per_column is None:
         y_shift_per_column = [0] * num_columns
 
-    assert len(y_shift_per_column) == num_columns, "y_shift_per_column must have the same length as num_columns"
+    assert len(y_shift_per_column) == num_columns, (
+        f"y_shift_per_column {len(y_shift_per_column)} must have " f"the same length as num_columns {num_columns}"
+    )
 
     positions = []
     for i in range(num_columns):
@@ -154,19 +165,22 @@ def generate_multi_columns_probe(
 
 
 def generate_linear_probe(
-    num_elec: int = 16, ypitch: float = 20, contact_shapes: str = "circle", contact_shape_params: dict = {"radius": 6}
+    num_elec: int = 16,
+    ypitch: float = 20,
+    contact_shapes: "circle" | "rect" | "square" = "circle",
+    contact_shape_params: dict = {"radius": 6},
 ) -> Probe:
     """Generate a one-column linear probe.
 
     Parameters
     ----------
-    num_elec : int
-        Number of electrodes, by default 16
-    ypitch : float
-        Pitch in y direction, by default 20
-    contact_shapes : str, default 'circle'
-        Shape of the contacts ('circle', 'rect', 'square')
-    contact_shape_params : dict, default {'radius': 6}
+    num_elec : int, default: 16
+        Number of electrodes
+    ypitch : float, default: 20
+        Pitch in y direction
+    contact_shapes : "circle" | "rect" | "square", default 'circle'
+        Shape of the contacts
+    contact_shape_params : dict, default: {'radius': 6}
         Parameters for the shape.
         For circle: {"radius": float}
         For square: {"width": float}
@@ -178,6 +192,10 @@ def generate_linear_probe(
         The generated probe
     """
 
+    assert (
+        _default_dict[contact_shapes] in contact_shape_params.keys()
+    ), "contact_shapes and contact_shape_params must be coordinated see docstring"
+
     probe = generate_multi_columns_probe(
         num_columns=1,
         num_contact_per_column=num_elec,
@@ -195,10 +213,10 @@ def generate_multi_shank(num_shank: int = 2, shank_pitch: list = [150, 0], **kar
 
     Parameters
     ----------
-    num_shank : int, default 2
+    num_shank : int, default: 2
         Number of shanks
-    shank_pitch : list, default [150,0]
-        Distance between shanks, by default [150, 0]
+    shank_pitch : list, default: [150,0]
+        Distance between shanks
 
     Returns
     -------

src/probeinterface/io.py

@@ -81,7 +81,10 @@ def write_probeinterface(file: str | Path, probe_or_probegroup: Probe | ProbeGro
     elif isinstance(probe_or_probegroup, ProbeGroup):
         probegroup = probe_or_probegroup
     else:
-        raise ValueError("write_probeinterface : needs a probe or probegroup")
+        raise TypeError(
+            f"write_probeinterface : needs a probe or probegroup you "
+            f"entered an object of type: {type(probe_or_probegroup)}"
+        )
 
     file = Path(file)
 
@@ -322,7 +325,9 @@ def write_BIDS_probe(folder: str | Path, probe_or_probegroup: Probe | ProbeGroup
     elif isinstance(probe_or_probegroup, ProbeGroup):
         probegroup = probe_or_probegroup
     else:
-        raise ValueError("probe_or_probegroup has to be" "of type Probe or ProbeGroup")
+        raise TypeError(
+            f"probe_or_probegroup has to be" "of type Probe or ProbeGroup " f"not type: {type(probe_or_probegroup)}"
+        )
     folder = Path(folder)
 
     # ensure that prefix and file type indicator are separated by an underscore
@@ -445,7 +450,7 @@ def read_prb(file: str | Path) -> ProbeGroup:
     """
 
     file = Path(file).absolute()
-    assert file.is_file()
+    assert file.is_file(), "'file given is not of type file"
     with file.open("r") as f:
         contents = f.read()
     contents = re.sub(r"range\(([\d,]*)\)", r"list(range(\1))", contents)

src/probeinterface/library.py

@@ -32,9 +32,9 @@ def download_probeinterface_file(manufacturer: str, probe_name: str):
 
     Parameters
     ----------
-    manufacturer : str
-        The probe manufacturer (e.g. 'cambridgeneurotech')
-    probe_name : str
+    manufacturer : "cambridgeneurotech" | "neuronexus"
+        The probe manufacturer
+    probe_name : str (see probeinterface_libary for options)
         The probe name
     """
     os.makedirs(cache_folder / manufacturer, exist_ok=True)
@@ -51,9 +51,9 @@ def get_from_cache(manufacturer: str, probe_name: str) -> Optional["Probe"]:
 
     Parameters
     ----------
-    manufacturer : str
-        The probe manufacturer (e.g. 'cambridgeneurotech', 'neuronexus')
-    probe_name : str
+    manufacturer : "cambridgeneurotech" | "neuronexus"
+        The probe manufacturer
+    probe_name : str (see probeinterface_libary for options)
         The probe name
 
     Returns
@@ -78,11 +78,11 @@ def get_probe(manufacturer: str, probe_name: str, name: Optional[str] = None) ->
 
     Parameters
     ----------
-    manufacturer : str
-        The probe manufacturer (e.g. 'cambridgeneurotech', 'neuronexus')
-    probe_name : str
+    manufacturer : "cambridgeneurotech" | "neuronexus"
+        The probe manufacturer
+    probe_name : str (see probeinterface_libary for options)
         The probe name
-    name : str or None
+    name : str | None, default: None
         Optional name for the probe
 
     Returns

src/probeinterface/plotting.py

@@ -4,7 +4,7 @@
 
 Depending on Probe.ndim, the plotting is done in 2D or 3D
 """
-
+from __future__ import annotations
 import numpy as np
 from matplotlib import path as mpl_path
 
@@ -13,18 +13,18 @@ def plot_probe(
     probe,
     ax=None,
     contacts_colors=None,
-    with_contact_id=False,
-    with_device_index=False,
-    text_on_contact=None,
-    contacts_values=None,
-    cmap="viridis",
-    title=True,
-    contacts_kargs={},
-    probe_shape_kwargs={},
-    xlims=None,
-    ylims=None,
-    zlims=None,
-    show_channel_on_click=False,
+    with_contact_id: bool = False,
+    with_device_index: bool = False,
+    text_on_contact: list | np.ndarray | None = None,
+    contacts_values: np.ndarray | None = None,
+    cmap: str = "viridis",
+    title: bool = True,
+    contacts_kargs: dict = {},
+    probe_shape_kwargs: dict = {},
+    xlims: tuple | None = None,
+    ylims: tuple | None = None,
+    zlims: tuple | None = None,
+    show_channel_on_click: bool = False,
 ):
     """Plot a Probe object.
     Generates a 2D or 3D axis, depending on Probe.ndim
@@ -33,34 +33,34 @@ def plot_probe(
     ----------
     probe : Probe
         The probe object
-    ax : matplotlib.axis, optional
-        The axis to plot the probe on. If None, an axis is created, by default None
-    contacts_colors : matplotlib color, optional
-        The color of the contacts, by default None
-    with_contact_id : bool, optional
-        If True, channel ids are displayed on top of the channels, by default False
-    with_device_index : bool, optional
-        If True, device channel indices are displayed on top of the channels, by default False
-    text_on_contact: None or list or numpy.array
+    ax : matplotlib.axis | None, default: None
+        The axis to plot the probe on. If None, an axis is created
+    contacts_colors : matplotlib color | None, default: None
+        The color of the contacts
+    with_contact_id : bool, default: False
+        If True, channel ids are displayed on top of the channels
+    with_device_index : bool, default: False
+        If True, device channel indices are displayed on top of the channels
+    text_on_contact: None | list | numpy.array, default: None
         Addintional text to plot on each contact
-    contacts_values : np.array, optional
-        Values to color the contacts with, by default None
-    cmap : str, optional
-        [description], by default 'viridis'
-    title : bool, optional
-        If True, the axis title is set to the probe name, by default True
-    contacts_kargs : dict, optional
-        Dict with kwargs for contacts (e.g. alpha, edgecolor, lw), by default {}
-    probe_shape_kwargs : dict, optional
-        Dict with kwargs for probe shape (e.g. alpha, edgecolor, lw), by default {}
-    xlims : tuple, optional
-        Limits for x dimension, by default None
-    ylims : tuple, optional
-        Limits for y dimension, by default None
-    zlims : tuple, optional
-        Limits for z dimension, by default None
-    show_channel_on_click : bool, optional
-        If True, the channel information is shown upon click, by default False
+    contacts_values : np.array, default: None
+        Values to color the contacts with
+    cmap : a colormap color, default: "viridis"
+         A colormap color
+    title : bool, default: True
+        If True, the axis title is set to the probe name
+    contacts_kargs : dict, default: {}
+        Dict with kwargs for contacts (e.g. alpha, edgecolor, lw)
+    probe_shape_kwargs : dict, default: {}
+        Dict with kwargs for probe shape (e.g. alpha, edgecolor, lw)
+    xlims : tuple | None, default: None
+        Limits for x dimension
+    ylims : tuple | None, default: None
+        Limits for y dimension
+    zlims : tuple | None, default: None
+        Limits for z dimension
+    show_channel_on_click : bool, default: False
+        If True, the channel information is shown upon click
 
     Returns
     -------
@@ -182,16 +182,18 @@ def on_press(event):
     return poly, poly_contour
 
 
-def plot_probe_group(probegroup, same_axes=True, **kargs):
+def plot_probe_group(probegroup, same_axes: bool = True, **kargs):
     """Plot all probes from a ProbeGroup
     Can be in an existing set of axes or separate axes.
 
     Parameters
     ----------
     probegroup : ProbeGroup
         The ProbeGroup to plot
-    same_axes : bool, optional
-        If True, the probes are plotted on the same axis, by default True
+    same_axes : bool, default: True
+        If True, the probes are plotted on the same axis
+    kargs: dict
+        see docstring for plot_probe for possible kargs
     """
 
     import matplotlib.pyplot as plt