From 432802b5c8aed9d642a2c2dd48d1bc7ce09dc459 Mon Sep 17 00:00:00 2001
From: Philipp Schlegel <observing@web.de>
Date: Mon, 18 Mar 2024 18:49:00 +0000
Subject: [PATCH] few fixes for find_mat_version: 1. Use default `dataset` 2.
 Deal with cases where none of the IDs are materialized 3. Use custom
 MaterializationMatchError error

---
 fafbseg/flywire/utils.py | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/fafbseg/flywire/utils.py b/fafbseg/flywire/utils.py
index 415faa0..f1c5998 100644
--- a/fafbseg/flywire/utils.py
+++ b/fafbseg/flywire/utils.py
@@ -607,11 +607,12 @@ def get_lr_position(x, coordinates='nm'):
     return (m[:, 0] - x[:, 0]) / 2
 
 
+@inject_dataset()
 def find_mat_version(ids,
                      verbose=True,
                      allow_multiple=False,
                      raise_missing=True,
-                     dataset='production'):
+                     dataset=None):
     """Find a materialization version (or live) for given IDs.
 
     Parameters
@@ -682,13 +683,13 @@ def find_mat_version(ids,
             print('Using live materialization')
         return 'live'
 
-    if allow_multiple:
+    if allow_multiple and any(latest_valid != 0):
         if all(latest_valid != 0):
             if verbose and not SILENCE_FIND_MAT_VERSION:
                 print(f"Found root IDs spread across {len(np.unique(latest_valid))} "
                       "materialization versions.")
             return latest_valid
-
+        
         msg = (f"Found root IDs spread across {len(np.unique(latest_valid)) - 1} "
                f"materialization versions but {(latest_valid == 0).sum()} IDs "
                "do not exist in any of the materialized tables.")
@@ -698,16 +699,18 @@ def find_mat_version(ids,
                 print(msg)
             return latest_valid
         else:
-            raise ValueError(msg)
+            raise MaterializationMatchError(msg)
 
     if dataset not in ('public, '):
-        raise ValueError('Given root IDs do not (co-)exist in any of the available '
-                        'materialization versions (including live). Try updating '
-                        'root IDs and rerun your query.')
+        raise MaterializationMatchError(
+            'Given root IDs do not (co-)exist in any of the available '
+            'materialization versions (including live). Try updating '
+            'root IDs and rerun your query.')
     else:
-        raise ValueError('Given root IDs do not (co-)exist in any of the available '
-                        'public materialization versions. Please make sure that '
-                        'the root IDs do exist and rerun your query.')
+        raise MaterializationMatchError(
+            'Given root IDs do not (co-)exist in any of the available '
+            'public materialization versions. Please make sure that '
+            'the root IDs do exist and rerun your query.')
 
 
 def _is_valid_version(ids, version, dataset):
@@ -751,3 +754,7 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         global SILENCE_FIND_MAT_VERSION
         SILENCE_FIND_MAT_VERSION = False
+
+
+class MaterializationMatchError(Exception):
+    pass