pyglove 0.5.0.dev202510230131__py3-none-any.whl → 0.5.0.dev202601060812__py3-none-any.whl

pyglove/core/io/file_system.py

@@ -14,8 +14,13 @@
 """Pluggable file system."""
 
 import abc
+import datetime
+import fnmatch
+import glob as std_glob
 import io
 import os
+import re
+import shutil
 from typing import Any, Literal, Optional, Union
 
 
@@ -84,6 +89,10 @@ class FileSystem(metaclass=abc.ABCMeta):
   def exists(self, path: Union[str, os.PathLike[str]]) -> bool:
     """Returns True if a path exists."""
 
+  @abc.abstractmethod
+  def glob(self, pattern: Union[str, os.PathLike[str]]) -> list[str]:
+    """Lists all files or sub-directories."""
+
   @abc.abstractmethod
   def listdir(self, path: Union[str, os.PathLike[str]]) -> list[str]:
     """Lists all files or sub-directories."""
@@ -92,6 +101,14 @@ class FileSystem(metaclass=abc.ABCMeta):
   def isdir(self, path: Union[str, os.PathLike[str]]) -> bool:
     """Returns True if a path is a directory."""
 
+  @abc.abstractmethod
+  def getctime(self, path: Union[str, os.PathLike[str]]) -> float:
+    """Returns the creation time of a file."""
+
+  @abc.abstractmethod
+  def getmtime(self, path: Union[str, os.PathLike[str]]) -> float:
+    """Returns the last modification time of a file."""
+
   @abc.abstractmethod
   def mkdir(
       self, path: Union[str, os.PathLike[str]], mode: int = 0o777
@@ -111,6 +128,14 @@ class FileSystem(metaclass=abc.ABCMeta):
   def rm(self, path: Union[str, os.PathLike[str]]) -> None:
     """Removes a file based on a path."""
 
+  @abc.abstractmethod
+  def rename(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    """Renames a file based on a path."""
+
   @abc.abstractmethod
   def rmdir(self, path: Union[str, os.PathLike[str]]) -> bool:
     """Removes a directory based on a path."""
@@ -119,6 +144,14 @@ class FileSystem(metaclass=abc.ABCMeta):
   def rmdirs(self, path: Union[str, os.PathLike[str]]) -> None:
     """Removes a directory chain based on a path."""
 
+  @abc.abstractmethod
+  def copy(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    """Copies a file to a new path."""
+
 
 def resolve_path(path: Union[str, os.PathLike[str]]) -> str:
   if isinstance(path, str):
@@ -137,9 +170,10 @@ def resolve_path(path: Union[str, os.PathLike[str]]) -> str:
 class StdFile(File):
   """The standard file."""
 
-  def __init__(self, file_object) -> None:
+  def __init__(self, file_object, path: Union[str, os.PathLike[str]]) -> None:
     super().__init__()
     self._file_object = file_object
+    self._path = path
 
   def read(self, size: Optional[int] = None) -> Union[str, bytes]:
     return self._file_object.read(size)
@@ -169,7 +203,7 @@ class StdFileSystem(FileSystem):
   def open(
       self, path: Union[str, os.PathLike[str]], mode: str = 'r', **kwargs
   ) -> File:
-    return StdFile(io.open(path, mode, **kwargs))
+    return StdFile(io.open(path, mode, **kwargs), path)
 
   def chmod(self, path: Union[str, os.PathLike[str]], mode: int) -> None:
     os.chmod(path, mode)
@@ -177,12 +211,21 @@ class StdFileSystem(FileSystem):
   def exists(self, path: Union[str, os.PathLike[str]]) -> bool:
     return os.path.exists(path)
 
+  def glob(self, pattern: Union[str, os.PathLike[str]]) -> list[str]:
+    return std_glob.glob(pattern)
+
   def listdir(self, path: Union[str, os.PathLike[str]]) -> list[str]:
     return os.listdir(path)
 
   def isdir(self, path: Union[str, os.PathLike[str]]) -> bool:
     return os.path.isdir(path)
 
+  def getctime(self, path: Union[str, os.PathLike[str]]) -> float:
+    return os.path.getctime(path)
+
+  def getmtime(self, path: Union[str, os.PathLike[str]]) -> float:
+    return os.path.getmtime(path)
+
   def mkdir(
       self, path: Union[str, os.PathLike[str]], mode: int = 0o777
   ) -> None:
@@ -196,6 +239,13 @@ class StdFileSystem(FileSystem):
   ) -> None:
     os.makedirs(path, mode, exist_ok)
 
+  def rename(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    os.rename(oldpath, newpath)
+
   def rm(self, path: Union[str, os.PathLike[str]]) -> None:
     os.remove(path)
 
@@ -205,6 +255,13 @@ class StdFileSystem(FileSystem):
   def rmdirs(self, path: Union[str, os.PathLike[str]]) -> None:
     os.removedirs(path)
 
+  def copy(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    shutil.copy(oldpath, newpath)
+
 
 #
 # Memory file system.
@@ -218,6 +275,9 @@ class MemoryFile(File):
     super().__init__()
     self._buffer = buffer
     self._pos = 0
+    now = datetime.datetime.now().timestamp()
+    self.ctime = now
+    self.mtime = now
 
   def read(self, size: Optional[int] = None) -> Union[str, bytes]:
     return self._buffer.read(size)
@@ -227,6 +287,7 @@ class MemoryFile(File):
 
   def write(self, content: Union[str, bytes]) -> None:
     self._buffer.write(content)
+    self.mtime = datetime.datetime.now().timestamp()
 
   def seek(self, offset: int, whence: Literal[0, 1, 2] = 0) -> int:
     return self._buffer.seek(offset, whence)
@@ -286,6 +347,21 @@ class MemoryFileSystem(FileSystem):
   def exists(self, path: Union[str, os.PathLike[str]]) -> bool:
     return self._locate(path) is not None
 
+  def glob(self, pattern: Union[str, os.PathLike[str]]) -> list[str]:
+    pattern = resolve_path(pattern)
+    regex = re.compile(fnmatch.translate(pattern))
+
+    results = []
+    def _traverse(node, prefix_parts):
+      for k, v in node.items():
+        path = self._prefix + '/'.join(prefix_parts + [k])
+        if regex.match(path):
+          results.append(path)
+        if isinstance(v, dict):
+          _traverse(v, prefix_parts + [k])
+    _traverse(self._root, [])
+    return results
+
   def listdir(self, path: Union[str, os.PathLike[str]]) -> list[str]:
     d = self._locate(path)
     if not isinstance(d, dict):
@@ -295,6 +371,24 @@ class MemoryFileSystem(FileSystem):
   def isdir(self, path: Union[str, os.PathLike[str]]) -> bool:
     return isinstance(self._locate(path), dict)
 
+  def getctime(self, path: Union[str, os.PathLike[str]]) -> float:
+    """Returns the creation time of a file."""
+    f = self._locate(path)
+    if isinstance(f, dict):
+      raise IsADirectoryError(path)
+    if f is None:
+      raise FileNotFoundError(path)
+    return f.ctime
+
+  def getmtime(self, path: Union[str, os.PathLike[str]]) -> float:
+    """Returns the last modification time of a file."""
+    f = self._locate(path)
+    if isinstance(f, dict):
+      raise IsADirectoryError(path)
+    if f is None:
+      raise FileNotFoundError(path)
+    return f.mtime
+
   def _parent_and_name(
       self, path: Union[str, os.PathLike[str]]
   ) -> tuple[dict[str, Any], str]:
@@ -338,6 +432,46 @@ class MemoryFileSystem(FileSystem):
         raise NotADirectoryError(path)
       current = entry
 
+  def rename(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    oldpath_str = resolve_path(oldpath)
+    newpath_str = resolve_path(newpath)
+
+    old_parent_dir, old_name = self._parent_and_name(oldpath_str)
+    entry = old_parent_dir.get(old_name)
+
+    if entry is None:
+      raise FileNotFoundError(oldpath_str)
+
+    new_entry = self._locate(newpath_str)
+    if new_entry is not None:
+      if isinstance(entry, dict):  # oldpath is dir
+        if not isinstance(new_entry, dict):
+          raise NotADirectoryError(
+              f'Cannot rename directory {oldpath_str!r} '
+              f'to non-directory {newpath_str!r}')
+        elif new_entry:
+          raise OSError(f'Directory not empty: {newpath_str!r}')
+        else:
+          self.rmdir(newpath_str)
+      else:  # oldpath is file
+        if isinstance(new_entry, dict):
+          raise IsADirectoryError(
+              f'Cannot rename non-directory {oldpath_str!r} '
+              f'to directory {newpath_str!r}')
+        else:
+          self.rm(newpath_str)
+    elif isinstance(entry, dict) and newpath_str.startswith(oldpath_str + '/'):
+      raise OSError(f'Cannot move directory {oldpath_str!r} '
+                    f'to a subdirectory of itself {newpath_str!r}')
+
+    new_parent_dir, new_name = self._parent_and_name(newpath_str)
+    new_parent_dir[new_name] = entry
+    del old_parent_dir[old_name]
+
   def rm(self, path: Union[str, os.PathLike[str]]) -> None:
     parent_dir, name = self._parent_and_name(path)
     entry = parent_dir.get(name)
@@ -379,6 +513,37 @@ class MemoryFileSystem(FileSystem):
       return not dir_dict
     _rmdir(self._root, self._internal_path(path))
 
+  def copy(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    old_file = self._locate(oldpath)
+    if old_file is None:
+      raise FileNotFoundError(oldpath)
+    if not isinstance(old_file, MemoryFile):
+      raise IsADirectoryError(oldpath)
+
+    if self.isdir(newpath):
+      _, old_file_name = self._parent_and_name(oldpath)
+      newpath = resolve_path(newpath)
+      newpath = newpath.rstrip('/') + '/' + old_file_name
+
+    # If newpath exists as file, remove it for overwrite.
+    if self.exists(newpath) and not self.isdir(newpath):
+      self.rm(newpath)
+    elif self.isdir(newpath):
+      raise IsADirectoryError(newpath)
+
+    old_pos = old_file.tell()
+    old_file.seek(0)
+    content = old_file.read()
+    old_file.seek(old_pos)
+
+    is_binary = isinstance(old_file._buffer, io.BytesIO)  # pylint: disable=protected-access
+    with self.open(newpath, 'wb' if is_binary else 'w') as f_new:
+      f_new.write(content)
+
 
 class _FileSystemRegistry:
   """File system registry."""
@@ -412,6 +577,260 @@ def add_file_system(prefix: str, fs: FileSystem) -> None:
 add_file_system('/mem/', MemoryFileSystem('/mem/'))
 
 
+try:
+  # pylint: disable=g-import-not-at-top
+  # pytype: disable=import-error
+  import fsspec
+  # pytype: enable=import-error
+  # pylint: enable=g-import-not-at-top
+except ImportError:
+  fsspec = None
+
+
+class FsspecFile(File):
+  """File object based on fsspec."""
+
+  def __init__(self, f):
+    self._f = f
+
+  def read(self, size: Optional[int] = None) -> Union[str, bytes]:
+    return self._f.read(size)
+
+  def readline(self) -> Union[str, bytes]:
+    return self._f.readline()
+
+  def write(self, content: Union[str, bytes]) -> None:
+    self._f.write(content)
+
+  def seek(self, offset: int, whence: Literal[0, 1, 2] = 0) -> int:
+    return self._f.seek(offset, whence)
+
+  def tell(self) -> int:
+    return self._f.tell()
+
+  def flush(self) -> None:
+    self._f.flush()
+
+  def close(self) -> None:
+    self._f.close()
+
+
+class FsspecFileSystem(FileSystem):
+  """File system based on fsspec."""
+
+  def open(
+      self, path: Union[str, os.PathLike[str]], mode: str = 'r', **kwargs
+  ) -> File:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    return FsspecFile(fsspec.open(path, mode, **kwargs).open())
+
+  def chmod(self, path: Union[str, os.PathLike[str]], mode: int) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    if hasattr(fs, 'chmod'):
+      fs.chmod(path, mode)
+
+  def exists(self, path: Union[str, os.PathLike[str]]) -> bool:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    return fs.exists(path)
+
+  def glob(self, pattern: Union[str, os.PathLike[str]]) -> list[str]:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(pattern)
+    protocol = fsspec.utils.get_protocol(pattern)
+    return [f'{protocol}:///{r.lstrip("/")}' for r in fs.glob(path)]
+
+  def listdir(self, path: Union[str, os.PathLike[str]]) -> list[str]:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    return [os.path.basename(f) for f in fs.ls(path, detail=False)]
+
+  def isdir(self, path: Union[str, os.PathLike[str]]) -> bool:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    return fs.isdir(path)
+
+  def getctime(self, path: Union[str, os.PathLike[str]]) -> float:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path_in_fs = fsspec.core.url_to_fs(path)
+    created = fs.created(path_in_fs)
+    if created is None:
+      raise OSError(f'c-time is not available for path {path!r}')
+    return created.timestamp()
+
+  def getmtime(self, path: Union[str, os.PathLike[str]]) -> float:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path_in_fs = fsspec.core.url_to_fs(path)
+    modified = fs.modified(path_in_fs)
+    if modified is None:
+      raise OSError(f'm-time is not available for path {path!r}')
+    return modified.timestamp()
+
+  def mkdir(
+      self,
+      path: Union[str, os.PathLike[str]], mode: int = 0o777
+  ) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    fs.mkdir(path)
+
+  def mkdirs(
+      self,
+      path: Union[str, os.PathLike[str]],
+      mode: int = 0o777,
+      exist_ok: bool = True,
+  ) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    fs.makedirs(path, exist_ok=exist_ok)
+
+  def rm(self, path: Union[str, os.PathLike[str]]) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    fs.rm(path)
+
+  def rename(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, old_path = fsspec.core.url_to_fs(oldpath)
+    fs2, new_path = fsspec.core.url_to_fs(newpath)
+    if fs.__class__ != fs2.__class__:
+      raise ValueError(
+          f'Rename across different filesystems is not supported: '
+          f'{type(fs)} vs {type(fs2)}'
+      )
+    fs.rename(old_path, new_path)
+
+  def rmdir(self, path: Union[str, os.PathLike[str]]) -> None:   # pytype: disable=signature-mismatch
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    fs.rmdir(path)
+
+  def rmdirs(self, path: Union[str, os.PathLike[str]]) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs, path = fsspec.core.url_to_fs(path)
+    fs.rm(path, recursive=True)
+
+  def copy(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    assert fsspec is not None, '`fsspec` is not installed.'
+    fs1, old_path_in_fs = fsspec.core.url_to_fs(oldpath)
+    fs2, new_path_in_fs = fsspec.core.url_to_fs(newpath)
+
+    if fs2.isdir(new_path_in_fs):
+      new_path_in_fs = os.path.join(
+          new_path_in_fs, os.path.basename(old_path_in_fs)
+      )
+
+    if fs1.__class__ == fs2.__class__:
+      fs1.copy(old_path_in_fs, new_path_in_fs)
+    else:
+      with fs1.open(old_path_in_fs, 'rb') as f1:
+        with fs2.open(new_path_in_fs, 'wb') as f2:
+          while True:
+            # Reads the file in 16MB chunks.
+            chunk = f1.read(16 * 1024 * 1024)
+            if not chunk:
+              break
+            f2.write(chunk)
+
+
+class _FsspecUriCatcher(FileSystem):
+  """File system to catch URI paths and redirect to FsspecFileSystem."""
+
+  # Catch all paths that contains '://' but not registered by
+  # available_protocols.
+  _URI_PATTERN = re.compile(r'^[a-zA-Z][a-zA-Z0-9+-.]*://.*')
+
+  def __init__(self):
+    super().__init__()
+    self._std_fs = StdFileSystem()
+    self._fsspec_fs = FsspecFileSystem()
+
+  def get_fs(self, path: Union[str, os.PathLike[str]]) -> FileSystem:
+    if self._URI_PATTERN.match(resolve_path(path)):
+      return self._fsspec_fs
+    return self._std_fs
+
+  def open(
+      self, path: Union[str, os.PathLike[str]], mode: str = 'r', **kwargs
+  ) -> File:
+    return self.get_fs(path).open(path, mode, **kwargs)
+
+  def chmod(self, path: Union[str, os.PathLike[str]], mode: int) -> None:
+    self.get_fs(path).chmod(path, mode)
+
+  def exists(self, path: Union[str, os.PathLike[str]]) -> bool:
+    return self.get_fs(path).exists(path)
+
+  def glob(self, pattern: Union[str, os.PathLike[str]]) -> list[str]:
+    return self.get_fs(pattern).glob(pattern)
+
+  def listdir(self, path: Union[str, os.PathLike[str]]) -> list[str]:
+    return self.get_fs(path).listdir(path)
+
+  def isdir(self, path: Union[str, os.PathLike[str]]) -> bool:
+    return self.get_fs(path).isdir(path)
+
+  def getctime(self, path: Union[str, os.PathLike[str]]) -> float:
+    return self.get_fs(path).getctime(path)
+
+  def getmtime(self, path: Union[str, os.PathLike[str]]) -> float:
+    return self.get_fs(path).getmtime(path)
+
+  def mkdir(
+      self,
+      path: Union[str, os.PathLike[str]],
+      mode: int = 0o777
+  ) -> None:
+    self.get_fs(path).mkdir(path, mode)
+
+  def mkdirs(
+      self,
+      path: Union[str, os.PathLike[str]],
+      mode: int = 0o777,
+      exist_ok: bool = True,
+  ) -> None:
+    self.get_fs(path).mkdirs(path, mode, exist_ok)
+
+  def rm(self, path: Union[str, os.PathLike[str]]) -> None:
+    self.get_fs(path).rm(path)
+
+  def rename(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    self.get_fs(oldpath).rename(oldpath, newpath)
+
+  def rmdir(self, path: Union[str, os.PathLike[str]]) -> None:   # pytype: disable=signature-mismatch
+    self.get_fs(path).rmdir(path)
+
+  def rmdirs(self, path: Union[str, os.PathLike[str]]) -> None:
+    self.get_fs(path).rmdirs(path)
+
+  def copy(
+      self,
+      oldpath: Union[str, os.PathLike[str]],
+      newpath: Union[str, os.PathLike[str]],
+  ) -> None:
+    self.get_fs(oldpath).copy(oldpath, newpath)
+
+
+if fsspec is not None:
+  fsspec_fs = FsspecFileSystem()
+  for p in fsspec.available_protocols():
+    add_file_system(p + '://', fsspec_fs)
+  add_file_system('', _FsspecUriCatcher())
+
+
 #
 # APIs for file IO.
 #
@@ -458,16 +877,37 @@ def writefile(
     chmod(path, perms)
 
 
+def rename(
+    oldpath: Union[str, os.PathLike[str]],
+    newpath: Union[str, os.PathLike[str]],
+) -> None:
+  """Renames a file."""
+  _fs.get(oldpath).rename(oldpath, newpath)
+
+
 def rm(path: Union[str, os.PathLike[str]]) -> None:
   """Removes a file."""
   _fs.get(path).rm(path)
 
 
+def copy(
+    oldpath: Union[str, os.PathLike[str]],
+    newpath: Union[str, os.PathLike[str]],
+) -> None:
+  """Copies a file."""
+  _fs.get(oldpath).copy(oldpath, newpath)
+
+
 def path_exists(path: Union[str, os.PathLike[str]]) -> bool:
   """Returns True if path exists."""
   return _fs.get(path).exists(path)
 
 
+def glob(pattern: Union[str, os.PathLike[str]]) -> list[str]:
+  """Lists all files or sub-directories."""
+  return _fs.get(pattern).glob(pattern)
+
+
 def listdir(
     path: Union[str, os.PathLike[str]], fullpath: bool = False
 ) -> list[str]:  # pylint: disable=redefined-builtin
@@ -483,6 +923,16 @@ def isdir(path: Union[str, os.PathLike[str]]) -> bool:
   return _fs.get(path).isdir(path)
 
 
+def getctime(path: Union[str, os.PathLike[str]]) -> float:
+  """Returns the creation time of a file."""
+  return _fs.get(path).getctime(path)
+
+
+def getmtime(path: Union[str, os.PathLike[str]]) -> float:
+  """Returns the last modification time of a file."""
+  return _fs.get(path).getmtime(path)
+
+
 def mkdir(path: Union[str, os.PathLike[str]], mode: int = 0o777) -> None:
   """Makes a directory."""
   _fs.get(path).mkdir(path, mode=mode)