pyglove 0.4.5.dev20240318__py3-none-any.whl → 0.4.5.dev202501132210__py3-none-any.whl

pyglove/core/hyper/evolvable.py

@@ -20,9 +20,9 @@ import types
 from typing import Any, Callable, List, Optional, Tuple, Union
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper import custom
 
 
@@ -44,7 +44,7 @@ class MutationPoint:
     parent: The parent node of the mutation point.
   """
   mutation_type: 'MutationType'
-  location: object_utils.KeyPath
+  location: utils.KeyPath
   old_value: Any
   parent: Optional[symbolic.Symbolic]
 
@@ -71,9 +71,9 @@ class Evolvable(custom.CustomHyper):
     mutation_points: List[MutationPoint] = []
     mutation_weights: List[float] = []
 
-    def _choose_mutation_point(k: object_utils.KeyPath,
-                               v: Any,
-                               p: Optional[symbolic.Symbolic]):
+    def _choose_mutation_point(
+        k: utils.KeyPath, v: Any, p: Optional[symbolic.Symbolic]
+    ):
       """Visiting function for a symbolic node."""
       def _add_point(mt: MutationType, k=k, v=v, p=p):
         mutation_points.append(MutationPoint(mt, k, v, p))
@@ -98,10 +98,9 @@ class Evolvable(custom.CustomHyper):
           reached_min_size = False
 
         for i, cv in enumerate(v):
-          ck = object_utils.KeyPath(i, parent=k)
+          ck = utils.KeyPath(i, parent=k)
           if not reached_max_size:
-            _add_point(MutationType.INSERT,
-                       k=ck, v=object_utils.MISSING_VALUE, p=v)
+            _add_point(MutationType.INSERT, k=ck, v=utils.MISSING_VALUE, p=v)
 
           if not reached_min_size:
             _add_point(MutationType.DELETE, k=ck, v=cv, p=v)
@@ -109,10 +108,12 @@ class Evolvable(custom.CustomHyper):
           # Replace type and value will be added in traverse.
           symbolic.traverse(cv, _choose_mutation_point, root_path=ck, parent=v)
           if not reached_max_size and i == len(v) - 1:
-            _add_point(MutationType.INSERT,
-                       k=object_utils.KeyPath(i + 1, parent=k),
-                       v=object_utils.MISSING_VALUE,
-                       p=v)
+            _add_point(
+                MutationType.INSERT,
+                k=utils.KeyPath(i + 1, parent=k),
+                v=utils.MISSING_VALUE,
+                p=v,
+            )
         return symbolic.TraverseAction.CONTINUE
       return symbolic.TraverseAction.ENTER
 
@@ -157,7 +158,7 @@ class Evolvable(custom.CustomHyper):
             point.location, point.old_value, point.parent)
     elif point.mutation_type == MutationType.INSERT:
       assert isinstance(point.parent, symbolic.List), point
-      assert point.old_value == object_utils.MISSING_VALUE, point
+      assert point.old_value == utils.MISSING_VALUE, point
       assert isinstance(point.location.key, int), point
       with symbolic.allow_writable_accessors():
         point.parent.insert(
@@ -175,24 +176,31 @@ class Evolvable(custom.CustomHyper):
 # We defer members declaration for Evolvable since the weights will reference
 # the definition of MutationType.
 symbolic.members([
-    ('initial_value', pg_typing.Object(symbolic.Symbolic),
-     'Symbolic value to involve.'),
-    ('node_transform', pg_typing.Callable(
-        [],
-        returns=pg_typing.Any()),
-     ''),
-    ('weights', pg_typing.Callable(
-        [
-            pg_typing.Object(MutationType),
-            pg_typing.Object(object_utils.KeyPath),
-            pg_typing.Any().noneable(),
-            pg_typing.Object(symbolic.Symbolic)
-        ], returns=pg_typing.Float(min_value=0.0)).noneable(),
-     ('An optional callable object that returns the unnormalized (e.g. '
-      'the sum of all probabilities do not have to sum to 1.0) mutation '
-      'probabilities for all the nodes in the symbolic tree, based on '
-      '(mutation type, location, old value, parent node). If None, all the '
-      'locations and mutation types will be sampled uniformly.')),
+    (
+        'initial_value',
+        pg_typing.Object(symbolic.Symbolic),
+        'Symbolic value to involve.',
+    ),
+    ('node_transform', pg_typing.Callable([], returns=pg_typing.Any()), ''),
+    (
+        'weights',
+        pg_typing.Callable(
+            [
+                pg_typing.Object(MutationType),
+                pg_typing.Object(utils.KeyPath),
+                pg_typing.Any().noneable(),
+                pg_typing.Object(symbolic.Symbolic),
+            ],
+            returns=pg_typing.Float(min_value=0.0),
+        ).noneable(),
+        (
+            'An optional callable object that returns the unnormalized (e.g.'
+            ' the sum of all probabilities do not have to sum to 1.0) mutation'
+            ' probabilities for all the nodes in the symbolic tree, based on'
+            ' (mutation type, location, old value, parent node). If None, all'
+            ' the locations and mutation types will be sampled uniformly.'
+        ),
+    ),
 ])(Evolvable)
 
 
@@ -200,25 +208,28 @@ def evolve(
     initial_value: symbolic.Symbolic,
     node_transform: Callable[
         [
-            object_utils.KeyPath,    # Location.
-            Any,                     # Old value.
-                                     # pg.MISSING_VALUE for insertion.
-            symbolic.Symbolic,       # Parent node.
+            utils.KeyPath,  # Location.
+            Any,  # Old value.
+            # pg.MISSING_VALUE for insertion.
+            symbolic.Symbolic,  # Parent node.
         ],
-        Any                          # Replacement.
+        Any,  # Replacement.
     ],
     *,
-    weights: Optional[Callable[
-        [
-            MutationType,  # Mutation type.
-            object_utils.KeyPath,    # Location.
-            Any,                     # Value.
-            symbolic.Symbolic,       # Parent.
-        ],
-        float                        # Mutation weight.
-    ]] = None,  # pylint: disable=bad-whitespace
+    weights: Optional[
+        Callable[
+            [
+                MutationType,  # Mutation type.
+                utils.KeyPath,  # Location.
+                Any,  # Value.
+                symbolic.Symbolic,  # Parent.
+            ],
+            float,  # Mutation weight.
+        ]
+    ] = None,  # pylint: disable=bad-whitespace
     name: Optional[str] = None,
-    hints: Optional[Any] = None) -> Evolvable:
+    hints: Optional[Any] = None
+) -> Evolvable:
   """An evolvable symbolic value.
 
   Example::

pyglove/core/hyper/numerical.py

@@ -17,9 +17,9 @@ import typing
 from typing import Any, Callable, Optional, Tuple
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper import base
 
 
@@ -62,8 +62,7 @@ class Float(base.HyperPrimitive):
           f'\'min_value\' must be positive when `scale` is {self.scale!r}. '
           f'encountered: {self.min_value}.')
 
-  def dna_spec(self,
-               location: Optional[object_utils.KeyPath] = None) -> geno.Float:
+  def dna_spec(self, location: Optional[utils.KeyPath] = None) -> geno.Float:
     """Returns corresponding DNASpec."""
     return geno.Float(
         min_value=self.min_value,
@@ -71,55 +70,74 @@ class Float(base.HyperPrimitive):
         scale=self.scale,
         hints=self.hints,
         name=self.name,
-        location=location or object_utils.KeyPath())
+        location=location or utils.KeyPath(),
+    )
 
   def _decode(self) -> float:
     """Decode a DNA into a float value."""
     dna = self._dna
     if not isinstance(dna.value, float):
       raise ValueError(
-          object_utils.message_on_path(
-              f'Expect float value. Encountered: {dna.value}.', self.sym_path))
+          utils.message_on_path(
+              f'Expect float value. Encountered: {dna.value}.', self.sym_path
+          )
+      )
     if dna.value < self.min_value:
       raise ValueError(
-          object_utils.message_on_path(
+          utils.message_on_path(
               f'DNA value should be no less than {self.min_value}. '
-              f'Encountered {dna.value}.', self.sym_path))
+              f'Encountered {dna.value}.',
+              self.sym_path,
+          )
+      )
 
     if dna.value > self.max_value:
       raise ValueError(
-          object_utils.message_on_path(
+          utils.message_on_path(
               f'DNA value should be no greater than {self.max_value}. '
-              f'Encountered {dna.value}.', self.sym_path))
+              f'Encountered {dna.value}.',
+              self.sym_path,
+          )
+      )
     return dna.value
 
   def encode(self, value: float) -> geno.DNA:
     """Encode a float value into a DNA."""
     if not isinstance(value, float):
       raise ValueError(
-          object_utils.message_on_path(
+          utils.message_on_path(
               f'Value should be float to be encoded for {self!r}. '
-              f'Encountered {value}.', self.sym_path))
+              f'Encountered {value}.',
+              self.sym_path,
+          )
+      )
     if value < self.min_value:
       raise ValueError(
-          object_utils.message_on_path(
+          utils.message_on_path(
               f'Value should be no less than {self.min_value}. '
-              f'Encountered {value}.', self.sym_path))
+              f'Encountered {value}.',
+              self.sym_path,
+          )
+      )
     if value > self.max_value:
       raise ValueError(
-          object_utils.message_on_path(
+          utils.message_on_path(
               f'Value should be no greater than {self.max_value}. '
-              f'Encountered {value}.', self.sym_path))
+              f'Encountered {value}.',
+              self.sym_path,
+          )
+      )
     return geno.DNA(value)
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool = False,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, 'Float']:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, 'Float']:
     """Validate candidates during value_spec binding time."""
     del allow_partial
     del child_transform
@@ -134,17 +152,23 @@ class Float(base.HyperPrimitive):
       if (float_spec.min_value is not None
           and self.min_value < float_spec.min_value):
         raise ValueError(
-            object_utils.message_on_path(
+            utils.message_on_path(
                 f'Float.min_value ({self.min_value}) should be no less than '
                 f'the min value ({float_spec.min_value}) of value spec: '
-                f'{float_spec}.', path))
+                f'{float_spec}.',
+                path,
+            )
+        )
       if (float_spec.max_value is not None
           and self.max_value > float_spec.max_value):
         raise ValueError(
-            object_utils.message_on_path(
+            utils.message_on_path(
                 f'Float.max_value ({self.max_value}) should be no greater than '
                 f'the max value ({float_spec.max_value}) of value spec: '
-                f'{float_spec}.', path))
+                f'{float_spec}.',
+                path,
+            )
+        )
     return (False, self)
 
   def is_leaf(self) -> bool:

pyglove/core/hyper/numerical_test.py

@@ -11,14 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for pyglove.hyper.Float."""
-
 import unittest
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper.numerical import Float
 from pyglove.core.hyper.numerical import floatv
 
@@ -44,12 +42,14 @@ class FloatTest(unittest.TestCase):
       floatv(-1.0, 1.0, 'log')
 
   def test_dna_spec(self):
-    self.assertTrue(symbolic.eq(
-        floatv(0.0, 1.0).dna_spec('a'),
-        geno.Float(
-            location=object_utils.KeyPath('a'),
-            min_value=0.0,
-            max_value=1.0)))
+    self.assertTrue(
+        symbolic.eq(
+            floatv(0.0, 1.0).dna_spec('a'),
+            geno.Float(
+                location=utils.KeyPath('a'), min_value=0.0, max_value=1.0
+            ),
+        )
+    )
 
   def test_decode(self):
     v = floatv(0.0, 1.0)

pyglove/core/hyper/object_template.py

@@ -16,14 +16,14 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper import base
 from pyglove.core.hyper import derived
 
 
-class ObjectTemplate(base.HyperValue, object_utils.Formattable):
+class ObjectTemplate(base.HyperValue, utils.Formattable):
   """Object template that encodes and decodes symbolic values.
 
   An object template can be created from a hyper value, which is a symbolic
@@ -131,18 +131,18 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
     """
     super().__init__()
     self._value = value
-    self._root_path = object_utils.KeyPath()
+    self._root_path = utils.KeyPath()
     self._compute_derived = compute_derived
     self._where = where
     self._parse_generators()
 
   @property
-  def root_path(self) -> object_utils.KeyPath:
+  def root_path(self) -> utils.KeyPath:
     """Returns root path."""
     return self._root_path
 
   @root_path.setter
-  def root_path(self, path: object_utils.KeyPath):
+  def root_path(self, path: utils.KeyPath):
     """Set root path."""
     self._root_path = path
 
@@ -150,7 +150,8 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
     """Parse generators from its templated value."""
     hyper_primitives = []
     def _extract_immediate_child_hyper_primitives(
-        path: object_utils.KeyPath, value: Any) -> bool:
+        path: utils.KeyPath, value: Any
+    ) -> bool:
       """Extract top-level hyper primitives."""
       if (isinstance(value, base.HyperValue)
           and (not self._where or self._where(value))):
@@ -162,13 +163,14 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
         hyper_primitives.append((path, value))
       elif isinstance(value, symbolic.Object):
         for k, v in value.sym_items():
-          object_utils.traverse(
-              v, _extract_immediate_child_hyper_primitives,
-              root_path=object_utils.KeyPath(k, path))
+          utils.traverse(
+              v,
+              _extract_immediate_child_hyper_primitives,
+              root_path=utils.KeyPath(k, path),
+          )
       return True
 
-    object_utils.traverse(
-        self._value, _extract_immediate_child_hyper_primitives)
+    utils.traverse(self._value, _extract_immediate_child_hyper_primitives)
     self._hyper_primitives = hyper_primitives
 
   @property
@@ -186,15 +188,15 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
     """Returns whether current template is constant value."""
     return not self._hyper_primitives
 
-  def dna_spec(
-      self, location: Optional[object_utils.KeyPath] = None) -> geno.Space:
+  def dna_spec(self, location: Optional[utils.KeyPath] = None) -> geno.Space:
     """Return DNA spec (geno.Space) from this template."""
     return geno.Space(
         elements=[
             primitive.dna_spec(primitive_location)
             for primitive_location, primitive in self._hyper_primitives
         ],
-        location=location or object_utils.KeyPath())
+        location=location or utils.KeyPath(),
+    )
 
   def _decode(self) -> Any:
     """Decode DNA into a value."""
@@ -202,9 +204,10 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
     assert dna is not None
     if not self._hyper_primitives and (dna.value is not None or dna.children):
       raise ValueError(
-          object_utils.message_on_path(
-              f'Encountered extra DNA value to decode: {dna!r}',
-              self._root_path))
+          utils.message_on_path(
+              f'Encountered extra DNA value to decode: {dna!r}', self._root_path
+          )
+      )
 
     # Compute hyper primitive values first.
     rebind_dict = {}
@@ -214,11 +217,14 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
     else:
       if len(dna.children) != len(self._hyper_primitives):
         raise ValueError(
-            object_utils.message_on_path(
+            utils.message_on_path(
                 f'The length of child values ({len(dna.children)}) is '
-                f'different from the number of hyper primitives '
+                'different from the number of hyper primitives '
                 f'({len(self._hyper_primitives)}) in ObjectTemplate. '
-                f'DNA={dna!r}, ObjectTemplate={self!r}.', self._root_path))
+                f'DNA={dna!r}, ObjectTemplate={self!r}.',
+                self._root_path,
+            )
+        )
       for i, (primitive_location, primitive) in enumerate(
           self._hyper_primitives):
         rebind_dict[primitive_location.path] = (
@@ -247,18 +253,18 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
       # TODO(daiyip): Currently derived value parsing is done at decode time,
       # which can be optimized by moving to template creation time.
       derived_values = []
-      def _extract_derived_values(
-          path: object_utils.KeyPath, value: Any) -> bool:
+      def _extract_derived_values(path: utils.KeyPath, value: Any) -> bool:
         """Extract top-level primitives."""
         if isinstance(value, derived.DerivedValue):
           derived_values.append((path, value))
         elif isinstance(value, symbolic.Object):
           for k, v in value.sym_items():
-            object_utils.traverse(
-                v, _extract_derived_values,
-                root_path=object_utils.KeyPath(k, path))
+            utils.traverse(
+                v, _extract_derived_values, root_path=utils.KeyPath(k, path)
+            )
         return True
-      object_utils.traverse(value, _extract_derived_values)
+
+      utils.traverse(value, _extract_derived_values)
 
       if derived_values:
         if not copied:
@@ -299,9 +305,9 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
       ValueError if value cannot be encoded by this template.
     """
     children = []
-    def _encode(path: object_utils.KeyPath,
-                template_value: Any,
-                input_value: Any) -> Any:
+    def _encode(
+        path: utils.KeyPath, template_value: Any, input_value: Any
+    ) -> Any:
       """Encode input value according to template value."""
       if (pg_typing.MISSING_VALUE == input_value
           and pg_typing.MISSING_VALUE != template_value):
@@ -339,10 +345,12 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
               f'TemplateOnlyKeys={template_keys - value_keys}, '
               f'InputOnlyKeys={value_keys - template_keys})')
         for key in template_value.sym_keys():
-          object_utils.merge_tree(
+          utils.merge_tree(
               template_value.sym_getattr(key),
               input_value.sym_getattr(key),
-              _encode, root_path=object_utils.KeyPath(key, path))
+              _encode,
+              root_path=utils.KeyPath(key, path),
+          )
       elif isinstance(template_value, symbolic.Dict):
         # Do nothing since merge will iterate all elements in dict and list.
         if not isinstance(input_value, dict):
@@ -358,19 +366,23 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
               f'value. (Path=\'{path}\', Template={template_value!r}, '
               f'Input={input_value!r})')
         for i, template_item in enumerate(template_value):
-          object_utils.merge_tree(
-              template_item, input_value[i], _encode,
-              root_path=object_utils.KeyPath(i, path))
+          utils.merge_tree(
+              template_item,
+              input_value[i],
+              _encode,
+              root_path=utils.KeyPath(i, path),
+          )
       else:
         if template_value != input_value:
           raise ValueError(
-              f'Unmatched value between template and input. '
-              f'(Path=\'{path}\', '
-              f'Template={object_utils.quote_if_str(template_value)}, '
-              f'Input={object_utils.quote_if_str(input_value)})')
+              'Unmatched value between template and input. '
+              f"(Path='{path}', "
+              f'Template={utils.quote_if_str(template_value)}, '
+              f'Input={utils.quote_if_str(input_value)})'
+          )
       return template_value
-    object_utils.merge_tree(
-        self._value, value, _encode, root_path=self._root_path)
+
+    utils.merge_tree(self._value, value, _encode, root_path=self._root_path)
     return geno.DNA(None, children)
 
   def try_encode(self, value: Any) -> Tuple[bool, geno.DNA]:
@@ -399,18 +411,18 @@ class ObjectTemplate(base.HyperValue, object_utils.Formattable):
              root_indent: int = 0,
              **kwargs) -> str:
     """Format this object."""
-    details = object_utils.format(
-        self._value, compact, verbose, root_indent, **kwargs)
+    details = utils.format(self._value, compact, verbose, root_indent, **kwargs)
     return f'{self.__class__.__name__}(value={details})'
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, 'ObjectTemplate']:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, 'ObjectTemplate']:
     """Validate candidates during value_spec binding time."""
     # Check if value_spec directly accepts `self`.
     if not value_spec.value_type or not isinstance(self, value_spec.value_type):

pyglove/core/io/__init__.py

@@ -16,5 +16,6 @@
 # pylint: disable=g-bad-import-order
 
 from pyglove.core.io.file_system import *
+from pyglove.core.io.sequence import *
 
 # pylint: enable=g-bad-import-order

pyglove/core/io/file_system.py

@@ -51,6 +51,10 @@ class File(metaclass=abc.ABCMeta):
   def tell(self) -> int:
     """Returns the current position of the file."""
 
+  @abc.abstractmethod
+  def flush(self) -> None:
+    """Flushes the written content to the storage."""
+
   @abc.abstractmethod
   def close(self) -> None:
     """Closes the file."""
@@ -112,7 +116,7 @@ class FileSystem(metaclass=abc.ABCMeta):
     """Removes a directory chain based on a path."""
 
 
-def _resolve_path(path: Union[str, os.PathLike[str]]) -> str:
+def resolve_path(path: Union[str, os.PathLike[str]]) -> str:
   if isinstance(path, str):
     return path
   elif hasattr(path, '__fspath__'):
@@ -148,6 +152,9 @@ class StdFile(File):
   def tell(self) -> int:
     return self._file_object.tell()
 
+  def flush(self) -> None:
+    self._file_object.flush()
+
   def close(self) -> None:
     self._file_object.close()
 
@@ -185,7 +192,7 @@ class StdFileSystem(FileSystem):
   def rm(self, path: Union[str, os.PathLike[str]]) -> None:
     os.remove(path)
 
-  def rmdir(self, path: Union[str, os.PathLike[str]]) -> None:
+  def rmdir(self, path: Union[str, os.PathLike[str]]) -> None:  # pytype: disable=signature-mismatch
     os.rmdir(path)
 
   def rmdirs(self, path: Union[str, os.PathLike[str]]) -> None:
@@ -220,6 +227,9 @@ class MemoryFile(File):
   def tell(self) -> int:
     return self._buffer.tell()
 
+  def flush(self) -> None:
+    pass
+
   def close(self) -> None:
     self.seek(0)
 
@@ -233,7 +243,7 @@ class MemoryFileSystem(FileSystem):
     self._prefix = prefix
 
   def _internal_path(self, path: Union[str, os.PathLike[str]]) -> str:
-    return '/' + _resolve_path(path).lstrip(self._prefix)
+    return '/' + resolve_path(path).lstrip(self._prefix)
 
   def _locate(self, path: Union[str, os.PathLike[str]]) -> Any:
     current = self._root
@@ -277,7 +287,7 @@ class MemoryFileSystem(FileSystem):
   def _parent_and_name(
       self, path: Union[str, os.PathLike[str]]
   ) -> tuple[dict[str, Any], str]:
-    path = _resolve_path(path)
+    path = resolve_path(path)
     rpos = path.rfind('/')
     assert rpos >= 0, path
     name = path[rpos + 1:]
@@ -326,7 +336,7 @@ class MemoryFileSystem(FileSystem):
       raise IsADirectoryError(path)
     del parent_dir[name]
 
-  def rmdir(self, path: Union[str, os.PathLike[str]]) -> None:
+  def rmdir(self, path: Union[str, os.PathLike[str]]) -> None:  # pytype: disable=signature-mismatch
     parent_dir, name = self._parent_and_name(path)
     entry = parent_dir.get(name)
     if entry is None:
@@ -372,7 +382,7 @@ class _FileSystemRegistry:
 
   def get(self, path: Union[str, os.PathLike[str]]) -> FileSystem:
     """Gets the file system for a path."""
-    path = _resolve_path(path)
+    path = resolve_path(path)
     for prefix, fs in self._filesystems:
       if path.startswith(prefix):
         return fs
@@ -475,4 +485,4 @@ def rmdir(path: Union[str, os.PathLike[str]]) -> bool:
 
 def rmdirs(path: Union[str, os.PathLike[str]]) -> bool:
   """Removes a directory chain until a parent directory is not empty."""
-  return _fs.get(path).rmdirs(path)
+  return _fs.get(path).rmdirs(path)  # pytype: disable=bad-return-type