From c9e595ab995c7a7f41c177b98494ed3009fa6c1f Mon Sep 17 00:00:00 2001
From: lare <lare@lare.cc>
Date: Wed, 11 Sep 2024 20:42:14 +0200
Subject: [PATCH] allow only one "type" to be requested by mntner

---
 __main__.py | 32 ++++++++++++++++++--------------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/__main__.py b/__main__.py
index a92e3c7..8b8cc2c 100644
--- a/__main__.py
+++ b/__main__.py
@@ -122,18 +122,19 @@ class Registry(object):
                 return True, self.index[objectType][objectFile]
         return False, "not found"
 
-    def get_all_by_mntner(self, mntner: str) -> tuple[bool, dict or str]:
-
+    def get_all_by_mntner(self, mntner: str, requestedType:str = None) -> tuple[bool, dict or str]:
+        "requestedType: string of type, if None returns all"
         def _load_by_mntner(mntner: str) -> dict:
             ret = {}
             for objectType in self.index:
-                for objectFile in self.index[objectType]:
-                    if mntner in self.index[objectType][objectFile]["mnt-by"]:
-                        if objectType in ret:
-                            ret[objectType][objectFile] = self.index[objectType][objectFile]
-                        else:
-                            ret[objectType] = {
-                                objectFile: self.index[objectType][objectFile]}
+                if requestedType == None or objectType == requestedType:
+                    for objectFile in self.index[objectType]:
+                        if mntner in self.index[objectType][objectFile]["mnt-by"]:
+                            if objectType in ret:
+                                ret[objectType][objectFile] = self.index[objectType][objectFile]
+                            else:
+                                ret[objectType] = {
+                                    objectFile: self.index[objectType][objectFile]}
             return ret
 
         # check if index is not yet initialized
@@ -141,22 +142,25 @@ class Registry(object):
             print("INFO: requested get_object, but index not yet created")
             self._build_index()
 
-        if mntner in self._cache["mntner-objects"]:
+        if mntner in self._cache["mntner-objects"] and ("None" in self._cache["mntner-objects"][mntner].keys() or requestedType in self._cache["mntner-objects"][mntner].keys()):
             if self._cache["expired"]:
                 # if the cache is expired: clear cache and reset "expired"
                 self._cache["mntner-objects"] = {}
                 self._cache["expired"] = False
-
                 ret = _load_by_mntner(mntner)
-                self._cache["mntner-objects"][mntner] = ret
+                self._cache["mntner-objects"][mntner] = {}
+                self._cache["mntner-objects"][mntner]["None" if requestedType == None else requestedType] = ret
                 if ret == {}:
                     return False, f"no objects found for {mntner}"
                 return True, ret
 
             else:
-                return True, self._cache["mntner-objects"][mntner]
+                return True, self._cache["mntner-objects"][mntner]["None" if requestedType == None else requestedType]
         else:
-            ret = _load_by_mntner(mntner)
+            if requestedType is not None and mntner in self._cache["mntner-objects"] and "None" in self._cache["mntner-objects"][mntner]:
+                ret = {requestedType: self._cache["mntner-objects"]}
+            else:
+                ret = _load_by_mntner(mntner)
             self._cache["mntner-objects"][mntner] = ret
             if ret == {}:
                 return False, f"no objects found for {mntner}"