Allow FormatterV2 in mock

lsst · Jul 12, 2024 · 78ff3a6 · 78ff3a6
1 parent 7da1557
commit 78ff3a6
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/lsst/pipe/base/tests/mocks/_storage_class.py b/python/lsst/pipe/base/tests/mocks/_storage_class.py
@@ -51,6 +51,7 @@
     DatasetType,
     Formatter,
     FormatterFactory,
+    FormatterV2,
     LookupKey,
     SerializedDatasetType,
     StorageClass,
@@ -606,7 +607,7 @@ def new_get_storage_class(self: StorageClassFactory, storageClassName: str) -> S
 
     def new_get_formatter_class_with_match(
         self: FormatterFactory, entity: Any
-    ) -> tuple[LookupKey, type[Formatter], dict[str, Any]]:
+    ) -> tuple[LookupKey, type[Formatter | FormatterV2], dict[str, Any]]:
         try:
             return original_get_formatter_class_with_match(self, entity)
         except KeyError:
@@ -629,7 +630,7 @@ def new_get_formatter_class_with_match(
 
     def new_get_formatter_with_match(
         self: FormatterFactory, entity: Any, *args: Any, **kwargs: Any
-    ) -> tuple[LookupKey, Formatter]:
+    ) -> tuple[LookupKey, Formatter | FormatterV2]:
         try:
             return original_get_formatter_with_match(self, entity, *args, **kwargs)
         except KeyError: