Remove dependency on TensorFlow in favor of using gpathlib/pathlib for internal and external file access, respectively.

PiperOrigin-RevId: 728310349

Author: timblakely

connectomics/common/file.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
+# Copyright 2025 The Google Research Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -26,9 +26,12 @@
 import dataclasses_json
 import tensorstore as ts
 
+PathLike = Union[str, pathlib.PurePath]
 
 T = TypeVar('T', bound=dataclasses_json.DataClassJsonMixin)
-PathLike = Union[str, pathlib.Path]
+
+# Local Alias
+Path = pathlib.Path
 
 
 def save_dataclass_json(
@@ -298,29 +301,6 @@ def load_dataclass_json(
   )
 
 
-# TODO(timblakely): Remove dependency on TF when there's a common API to read
-# files internally and externally.
-import tensorflow as tf
-
-Copy = tf.io.gfile.copy
-DeleteRecursively = tf.io.gfile.rmtree
-Exists = tf.io.gfile.exists
-Glob = tf.io.gfile.glob
-IsDirectory = tf.io.gfile.isdir
-ListDirectory = tf.io.gfile.listdir
-MakeDirs = tf.io.gfile.makedirs
-MkDir = tf.io.gfile.mkdir
-Open = tf.io.gfile.GFile
-Remove = tf.io.gfile.remove
-Rename = tf.io.gfile.rename
-Stat = tf.io.gfile.stat
-Walk = tf.io.gfile.walk
-
-GFile = tf.io.gfile.GFile
-
-NotFoundError = tf.errors.NotFoundError
-
-
 # TODO(timblakely): Remove in favor the above serialization.
 def load_dataclass(
     constructor: type[T], v: Union[str, dict[str, Any], T, None]
@@ -344,7 +324,7 @@ def load_dataclass(
       return constructor.from_json(v)
     except json.JSONDecodeError:
       # File path; attempt to load.
-      with Open(v) as f:
+      with Path(v).open() as f:
         return constructor.from_json(f.read())
   else:
     return constructor.from_dict(typing.cast(dict[str, Any], v))

connectomics/common/file_test.py

@@ -134,7 +134,7 @@ def test_dataclass_from_json(self):
 
   def test_dataclass_from_file(self):
     fname = os.path.join(FLAGS.test_tmpdir, 'dc_file')
-    with file.Open(fname, 'w') as f:
+    with file.Path(fname).open('wt') as f:
       f.write("""{
         "spec": {
           "test": "foo"
@@ -179,7 +179,7 @@ def test_dataclass_loader(self):
     self.assertEqual(new_conf, ts_conf)
 
     fname = os.path.join(FLAGS.test_tmpdir, 'dc_loader')
-    with file.Open(fname, 'w') as f:
+    with file.Path(fname).open('wt') as f:
       f.write(ts_conf.to_json())
     new_conf = loader(fname)
     self.assertEqual(new_conf, ts_conf)

connectomics/common/io_utils.py

@@ -44,7 +44,7 @@ def load_equivalences(
   equiv_graph = nx.Graph()
 
   for path in paths:
-    with file.Open(path, "r") as f:
+    with file.Path(path).open("rt") as f:
       reader = pd.read_csv(
           f, sep=",", engine="c", comment="#", chunksize=4096, header=None
       )
@@ -64,13 +64,13 @@ def load_equivalences(
 
 def load_relabel_map(path: str) -> dict[int, int]:
   """Loads a label map from a text file."""
-  with file.Open(path) as f:
+  with file.Path(path).open("r") as f:
     df = pd.read_csv(f, engine="c", comment="#", header=None, dtype=np.uint64)
     return {int(a): int(b) for a, b in df.to_numpy(dtype=np.uint64)}
 
 
 def load_object_list(path: str) -> set[int]:
   """Loads an object list from a text file."""
-  with file.Open(path) as f:
+  with file.Path(path).open("r") as f:
     df = pd.read_csv(f, engine="c", comment="#", header=None, dtype=np.uint64)
     return set(int(x) for x in df.to_numpy(dtype=np.uint64).ravel())

connectomics/jax/models/util.py

@@ -187,11 +187,11 @@ def model_from_dict_config(
 
 def save_config(config: ml_collections.ConfigDict, path: file.PathLike):
   """Saves model config to a file."""
-  with file.Open(path, 'w') as f:
+  with file.Path(path).open('wt') as f:
     f.write(config.to_json_best_effort() + '\n')
 
 
 def load_config(path: file.PathLike) -> ml_collections.ConfigDict:
   """Loads a model config from a file."""
-  with file.Open(path, 'r') as f:
+  with file.Path(path).open('rt') as f:
     return ml_collections.ConfigDict(json.loads(f.read()))

connectomics/volume/decorators.py

@@ -1576,7 +1576,7 @@ def __init__(self, dim: int | str, json_path: str):
     self._json_path = json_path
 
   def decorate(self, input_ts: ts.TensorStore) -> ts.TensorStore:
-    indices = json_lib.loads(file.Open(self._json_path, 'r').read())
+    indices = json_lib.loads(file.Path(self._json_path).read_text())
     return input_ts[ts.d[self._dim][indices]]