Clean up docs and add list_all_probes function

Author: alejoe91

doc/library.rst

@@ -1,27 +1,58 @@
 Probeinterface public library
 =============================
 
-Probeinterface also handles a collection of probe descriptions on the
-`GitHub platform <https://github.com/SpikeInterface/probeinterface_library>`_
+Probeinterface also handles a collection of probe descriptions in the
+`ProbeInterface library <https://github.com/SpikeInterface/probeinterface_library>`_
 
-The python module has a simple function to download and cache locally by using `get_probe(...)` ::
+The python module has a simple function to download and cache locally by using ``get_probe(...)``:
+
+
+.. code-block:: python
 
     from probeinterface import get_probe
-    probe = get_probe(manufacturer='neuronexus',
-                probe_name='A1x32-Poly3-10mm-50-177')
+    probe = get_probe(
+        manufacturer='neuronexus',
+        probe_name='A1x32-Poly3-10mm-50-177'
+    )
+
+
+Once a probe is downloaded, it is cached locally for future use.
+
+There are several helper functions to explore the library:
+
+.. code-block:: python
+
+    from probeinterface.library import (
+        list_manufacturers,
+        list_probes_by_manufacturer,
+        list_all_probes
+    )
+
+    # List all manufacturers
+    manufacturers = list_manufacturers()
+
+    # List all probes for a given manufacturer
+    probes = list_probes_by_manufacturer('neuronexus')
+
+    # List all probes in the library
+    all_probes = list_all_probes()
+
+    # Cache all probes locally
+    cache_full_library()
+
 
+Each function has an optional ``tag`` argument to specify a git tag/branch/commit to get a specific version of the library.
 
-We expect to build rapidly commonly used probes in this public repository.
 
-How to contribute
------------------
+How to contribute to the library
+--------------------------------
 
-TODO: explain with more details
+Each probe in the library is represented by a JSON file and an image.
+To contribute a new probe to the library, follow these steps:
 
-  1. Generate the JSON file with probeinterface (or directly
-      with another language)
+  1. Generate the JSON file with probeinterface (or directly with another language)
   2. Generate an image of the probe with the `plot_probe` function in probeinterface
   3. Clone the `probeinterface_library repo <https://github.com/SpikeInterface/probeinterface_library>`_
-  4. Put the JSON file and image into the correct folder or make a new folder (following the format of the repo)
+  4. Put the JSON file and image into the correct folder: ``probeinterface_library/<manufacturer>/<model_name>/```
   5. Push to one of your branches with a git client
   6. Make a pull request to the main repo

src/probeinterface/__init__.py

@@ -41,9 +41,11 @@
 )
 from .library import (
     get_probe,
-    list_manufacturers_in_library,
-    list_probes_in_library,
+    list_manufacturers,
+    list_probes_by_manufacturer,
+    list_all_probes,
     get_tags_in_library,
     cache_full_library,
+    clear_cache
 )
 from .wiring import get_available_pathways

src/probeinterface/library.py

@@ -146,22 +146,44 @@ def get_probe(
     return probe
 
 
-def cache_full_library(tag=None) -> None:
+def cache_full_library(tag=None) -> None:  # pragma: no cover
     """
     Download all probes from the library to the cache directory.
     """
-    manufacturers = list_manufacturers_in_library(tag=tag)
+    manufacturers = list_manufacturers(tag=tag)
 
     for manufacturer in manufacturers:
-        probes = list_probes_in_library(manufacturer, tag=tag)
+        probes = list_probes_by_manufacturer(manufacturer, tag=tag)
         for probe_name in probes:
             try:
                 download_probeinterface_file(manufacturer, probe_name, tag=tag)
             except Exception as e:
                 warnings.warn(f"Could not download {manufacturer}/{probe_name} (tag: {tag}): {e}")
 
 
-def list_manufacturers_in_library(tag=None) -> list[str]:
+def clear_cache(tag=None) -> None:  # pragma: no cover
+    """
+    Clear the cache folder for probeinterface library files.
+
+    Parameters
+    ----------
+    tag : str | None, default: None
+        Optional tag for the probe
+    """
+    cache_folder = get_cache_folder()
+    if tag is not None:
+        cache_folder_tag = cache_folder / tag
+        if cache_folder_tag.is_dir():
+            import shutil
+
+            shutil.rmtree(cache_folder_tag)
+    else:
+        import shutil
+
+        shutil.rmtree(cache_folder)
+
+
+def list_manufacturers(tag=None) -> list[str]:
     """
     Get the list of available manufacturers in the library
 
@@ -173,7 +195,7 @@ def list_manufacturers_in_library(tag=None) -> list[str]:
     return list_github_folders("SpikeInterface", "probeinterface_library", ref=tag)
 
 
-def list_probes_in_library(manufacturer: str, tag=None) -> list[str]:
+def list_probes_by_manufacturer(manufacturer: str, tag=None) -> list[str]:
     """
     Get the list of available probes for a given manufacturer
 
@@ -190,6 +212,23 @@ def list_probes_in_library(manufacturer: str, tag=None) -> list[str]:
     return list_github_folders("SpikeInterface", "probeinterface_library", path=manufacturer, ref=tag)
 
 
+def list_all_probes(tag=None) -> dict[str, list[str]]:
+    """
+    Get the list of all available probes in the library
+
+    Returns
+    -------
+    all_probes : dict
+        Dictionary with manufacturers as keys and list of probes as values
+    """
+    all_probes = {}
+    manufacturers = list_manufacturers(tag=tag)
+    for manufacturer in manufacturers:
+        probes = list_probes_by_manufacturer(manufacturer, tag=tag)
+        all_probes[manufacturer] = probes
+    return all_probes
+
+
 def get_tags_in_library() -> list[str]:
     """
     Get the list of available tags in the library

tests/test_library.py

@@ -1,13 +1,18 @@
 import os
+import pytest
+
 from probeinterface import Probe
 from probeinterface.library import (
     download_probeinterface_file,
     get_from_cache,
     get_probe,
     get_tags_in_library,
-    list_manufacturers_in_library,
-    list_probes_in_library,
+    list_manufacturers,
+    list_probes_by_manufacturer,
+    list_all_probes,
     get_cache_folder,
+    cache_full_library,
+    clear_cache
 )
 
 
@@ -77,24 +82,44 @@ def test_available_tags():
             assert len(tag) > 0
 
 
-def test_list_manufacturers_in_library():
-    manufacturers = list_manufacturers_in_library()
+def test_list_manufacturers():
+    manufacturers = list_manufacturers()
     assert isinstance(manufacturers, list)
     assert "neuronexus" in manufacturers
     assert "imec" in manufacturers
 
 
-def test_list_probes_in_library():
-    manufacturers = list_manufacturers_in_library()
+def test_list_probes():
+    manufacturers = list_all_probes()
     for manufacturer in manufacturers:
-        probes = list_probes_in_library(manufacturer)
+        probes = list_probes_by_manufacturer(manufacturer)
         assert isinstance(probes, list)
         assert len(probes) > 0
 
 
+@pytest.mark.skip(reason="long test that downloads the full library")
+def test_cache_full_library():
+    tag = get_tags_in_library()[0] if len(get_tags_in_library()) > 0 else None
+    print(tag)
+    cache_full_library(tag=tag)
+    all_probes = list_all_probes(tag=tag)
+    # spot check that a known probe is in the cache
+    for manufacturer, probes in all_probes.items():
+        for probe_name in probes:
+            probe = get_from_cache(manufacturer, probe_name, tag=tag)
+            assert isinstance(probe, Probe)
+
+    clear_cache(tag=tag)
+    for manufacturer, probes in all_probes.items():
+        for probe_name in probes:
+            probe = get_from_cache(manufacturer, probe_name, tag=tag)
+            assert probe is None
+
+
 if __name__ == "__main__":
     test_download_probeinterface_file()
     test_get_from_cache()
     test_get_probe()
-    test_list_manufacturers_in_library()
-    test_list_probes_in_library()
+    test_list_manufacturers()
+    test_list_probes()
+    test_cache_full_library()