From 04043bcb4a0339e7b7226b9fceba8c4cab0698f4 Mon Sep 17 00:00:00 2001
From: Adam Tyson <code@adamltyson.com>
Date: Tue, 29 Oct 2024 17:10:03 +0000
Subject: [PATCH] Make list atlas function more robust (#417)

* Make list atlas function more robust

* Add tests
---
 brainglobe_atlasapi/list_atlases.py | 16 +++++++++++-----
 tests/atlasapi/test_list_atlases.py | 10 ++++++++--
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/brainglobe_atlasapi/list_atlases.py b/brainglobe_atlasapi/list_atlases.py
index 9a5da303..40ea68ce 100644
--- a/brainglobe_atlasapi/list_atlases.py
+++ b/brainglobe_atlasapi/list_atlases.py
@@ -2,6 +2,8 @@
     Some functionality to list all available and downloaded brainglobe atlases
 """
 
+import re
+
 from rich import print as rprint
 from rich.panel import Panel
 from rich.table import Table
@@ -43,11 +45,15 @@ def get_local_atlas_version(atlas_name):
     """
 
     brainglobe_dir = config.get_brainglobe_dir()
-    return [
-        f.name.split("_v")[1]
-        for f in brainglobe_dir.glob(f"*{atlas_name}*")
-        if f.is_dir()
-    ][0]
+    try:
+        return [
+            re.search(r"_v(\d+\.\d+)$", f.name).group(1)
+            for f in brainglobe_dir.glob(f"*{atlas_name}*")
+            if f.is_dir() and re.search(r"_v(\d+\.\d+)$", f.name)
+        ][0]
+    except IndexError:
+        print(f"No atlas found with the name: {atlas_name}")
+        return None
 
 
 def get_all_atlases_lastversions():
diff --git a/tests/atlasapi/test_list_atlases.py b/tests/atlasapi/test_list_atlases.py
index 945357e6..b32f24b6 100644
--- a/tests/atlasapi/test_list_atlases.py
+++ b/tests/atlasapi/test_list_atlases.py
@@ -17,12 +17,18 @@ def test_get_downloaded_atlases():
     assert "example_mouse_100um" in available_atlases
 
 
-def test_get_local_atlas_version():
+def test_get_local_atlas_version_real_atlas():
     v = get_local_atlas_version("example_mouse_100um")
-
     assert len(v.split(".")) == 2
 
 
+def test_get_local_atlas_version_missing_atlas(capsys):
+    atlas_name = "unicorn_atlas"
+    assert get_local_atlas_version(atlas_name) is None
+    captured = capsys.readouterr()
+    assert f"No atlas found with the name: {atlas_name}" in captured.out
+
+
 def test_lastversions():
     last_versions = get_atlases_lastversions()
     example_atlas = last_versions["example_mouse_100um"]