diff --git a/src/cardiac_geometries/aha.py b/src/cardiac_geometries/aha.py
index 9626871..401acbf 100644
--- a/src/cardiac_geometries/aha.py
+++ b/src/cardiac_geometries/aha.py
@@ -30,7 +30,8 @@ def cartesian_to_prolate_ellipsoidal(x, y, z, a):
 
 def get_level(regions, mu):
     A = np.intersect1d(
-        np.where((regions.T[3] <= mu))[0], np.where((mu <= regions.T[0]))[0]
+        np.where((regions.T[3] <= mu))[0],
+        np.where((mu <= regions.T[0]))[0],
     )
     if len(A) == 0:
         return [np.shape(regions)[0] + 1]
@@ -120,7 +121,10 @@ def value_shape(self):
 
 
 def lv_aha(
-    geometry: Geometry, r_long_endo: float, r_short_endo: float, mu_base: float
+    geometry: Geometry,
+    r_long_endo: float,
+    r_short_endo: float,
+    mu_base: float,
 ) -> Geometry:
     foc = focal(r_long_endo=r_long_endo, r_short_endo=r_short_endo)
 
diff --git a/src/cardiac_geometries/geometry.py b/src/cardiac_geometries/geometry.py
index 2973a2c..51e6235 100644
--- a/src/cardiac_geometries/geometry.py
+++ b/src/cardiac_geometries/geometry.py
@@ -112,6 +112,19 @@ def to_dict(self):
         return self._asdict()
 
 
+def find_schema_in_folder(folder: Path) -> Optional[Dict[str, H5Path]]:
+    # First look for a schema file
+    for f in folder.iterdir():
+        if f.suffix == ".json":
+            try:
+                schema = load_schema(f)
+            except Exception:
+                pass
+            else:
+                return schema
+    return None
+
+
 def load_schema(path: Path) -> Optional[Dict[str, H5Path]]:
     if not path.is_file():
         return None
@@ -174,7 +187,10 @@ def read(
         raise RuntimeError(f"Unknown file format for {fname}")
 
 
-def extract_fname_group(fname: str, folder=".") -> Tuple[Path, Optional[str]]:
+def extract_fname_group(
+    fname: str,
+    folder: Union[Path, str] = ".",
+) -> Tuple[Path, Optional[str]]:
     fg = fname.split(":")
     if len(fg) == 1:
         return Path(folder) / fg[0], None
@@ -466,24 +482,13 @@ def from_folder(cls, folder, schema: Optional[Dict[str, H5Path]] = None):
         folder = Path(folder)
 
         if schema is None:
-            # First look for a schema file
-            for f in folder.iterdir():
-                if f.suffix == ".json":
-                    try:
-                        schema = load_schema(f)
-                    except Exception:
-                        continue
-                    else:
-                        break
-
-            else:
+            if (schema := find_schema_in_folder(folder)) is None:
                 schema = cls.default_schema()
-
+        assert schema is not None
         # Load mesh first
         data = {}
 
         for name, p in schema.items():
-            print(name, p)
             if p.fname == "":
                 continue
             if not p.is_mesh: