From 177081809ca57d5e5e8d8602ec41a14fa31068fd Mon Sep 17 00:00:00 2001
From: Ian Davies <iandavies@google.com>
Date: Wed, 13 Nov 2024 04:18:49 -0800
Subject: [PATCH] Allow passing either file bytes or filename when loading .mat
 files.

PiperOrigin-RevId: 696075023
---
 torax/geometry_loader.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/torax/geometry_loader.py b/torax/geometry_loader.py
index f8e55e31..2ae5dc93 100644
--- a/torax/geometry_loader.py
+++ b/torax/geometry_loader.py
@@ -13,8 +13,10 @@
 # limitations under the License.
 
 """File I/O for loading geometry files."""
+
 import enum
 import os
+from typing import IO
 
 import eqdsk
 import numpy as np
@@ -56,8 +58,8 @@ def _load_CHEASE_data(  # pylint: disable=invalid-name
   }
 
 
-def _load_fbt_data(file_path: str) -> dict[str, np.ndarray]:
-  """Loads the data from a FBT-LY file into a dictionary."""
+def _load_fbt_data(file_path: str | IO[bytes]) -> dict[str, np.ndarray]:
+  """Loads data into a dictionary from an FBT-LY file or file path."""
   return scipy.io.loadmat(file_path, squeeze_me=True)