pyglove 0.4.5.dev20240319__py3-none-any.whl → 0.4.5.dev202501140808__py3-none-any.whl

pyglove/core/{object_utils → utils}/json_conversion.py

@@ -16,13 +16,14 @@
 import abc
 import base64
 import collections
+import contextlib
 import importlib
 import inspect
 import marshal
 import pickle
 import types
 import typing
-from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Type, TypeVar, Union
+from typing import Any, Callable, ContextManager, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Type, TypeVar, Union
 
 # Nestable[T] is a (maybe) nested structure of T, which could be T, a Dict
 # a List or a Tuple of Nestable[T]. We use a Union to fool PyType checker to
@@ -36,7 +37,7 @@ JSONPrimitiveType = Union[int, float, bool, str]
 # pytype doesn't support recursion. Use Any instead of 'JSONValueType'
 # in List and Dict.
 JSONListType = List[Any]
-JSONDictType = Dict[str, Any]
+JSONDictType = Dict[Union[str, int], Any]
 JSONValueType = Union[JSONPrimitiveType, JSONListType, JSONDictType]
 
 # pylint: enable=invalid-name
@@ -57,6 +58,7 @@ class _TypeRegistry:
     # registered for a user class.
     self._type_to_cls_map = dict()
     self._prefix_mapping = dict()
+    self._ondemand_registry_stack = []
 
   def register(
       self, type_name: str, cls: Type[Any], override_existing: bool = False
@@ -89,9 +91,32 @@ class _TypeRegistry:
     """Returns whether a type name is registered."""
     return type_name in self._type_to_cls_map
 
+  @contextlib.contextmanager
+  def load_types_for_deserialization(
+      self,
+      *types_to_deserialize: Type[Any],
+  ) -> Iterator[Dict[str, Type[Any]]]:
+    """Context manager for loading unregistered types for deserialization."""
+    if self._ondemand_registry_stack:
+      stack_top = dict(self._ondemand_registry_stack[-1])
+    else:
+      stack_top = {}
+    stack_top.update({t.__name__: t for t in types_to_deserialize})
+    try:
+      self._ondemand_registry_stack.append(stack_top)
+      yield stack_top
+    finally:
+      self._ondemand_registry_stack.pop()
+
   def class_from_typename(
       self, type_name: str) -> Optional[Type[Any]]:
     """Get class from type name."""
+    if self._ondemand_registry_stack:
+      top_registry = self._ondemand_registry_stack[-1]
+      class_name = type_name.split('.')[-1]
+      if class_name in top_registry:
+        return top_registry[class_name]
+
     cls = self._type_to_cls_map.get(type_name, None)
     if cls is None:
       # Modules could be renamed, to load legacy serialized objects, we
@@ -177,9 +202,8 @@ class JSONConvertible(metaclass=abc.ABCMeta):
     Returns:
       An instance of cls.
     """
-    del kwargs
     assert isinstance(json_value, dict)
-    init_args = {k: from_json(v) for k, v in json_value.items()
+    init_args = {k: from_json(v, **kwargs) for k, v in json_value.items()
                  if k != JSONConvertible.TYPE_NAME_KEY}
     return cls(**init_args)
 
@@ -249,6 +273,36 @@ class JSONConvertible(metaclass=abc.ABCMeta):
     """Returns an iterator of registered (serialization key, class) tuples."""
     return cls._TYPE_REGISTRY.iteritems()
 
+  @classmethod
+  def load_types_for_deserialization(
+      cls,
+      *types_to_deserialize: Type[Any]
+      ) -> ContextManager[Dict[str, Type[Any]]]:
+    """Context manager for loading unregistered types for deserialization.
+
+    Example::
+
+      class A(pg.Object):
+        auto_register = False
+        x: int
+
+      class B(A):
+        y: str
+      with pg.JSONConvertile.load_types_for_deserialization(A, B):
+          pg.from_json_str(A(1).to_json_str())
+          pg.from_json_str(B(1, 'hi').to_json_str())
+
+    Args:
+      *types_to_deserialize: A list of types to be loaded for deserialization.
+
+    Returns:
+      A context manager within which the objects of the requested types
+        could be deserialized.
+    """
+    return cls._TYPE_REGISTRY.load_types_for_deserialization(
+        *types_to_deserialize
+    )
+
   @classmethod
   def to_json_dict(
       cls,
@@ -355,54 +409,68 @@ def to_json(value: Any, **kwargs) -> Any:
     JSON value.
   """
   if isinstance(value, (type(None), bool, int, float, str)):
-    return value
+    v = value
   elif isinstance(value, JSONConvertible):
-    return value.to_json(**kwargs)
+    v = value.to_json(**kwargs)
   elif isinstance(value, tuple):
-    return [JSONConvertible.TUPLE_MARKER] + to_json(list(value), **kwargs)
+    v = [JSONConvertible.TUPLE_MARKER] + to_json(list(value), **kwargs)
   elif isinstance(value, list):
-    return [to_json(item, **kwargs) for item in value]
+    v = [to_json(item, **kwargs) for item in value]
   elif isinstance(value, dict):
-    return {k: to_json(v, **kwargs) for k, v in value.items()}
+    v = {k: to_json(v, **kwargs) for k, v in value.items()}
   elif isinstance(value, (type, typing.GenericAlias)):  # pytype: disable=module-attr
-    return _type_to_json(value)
+    v = _type_to_json(value)
   elif inspect.isbuiltin(value):
-    return _builtin_function_to_json(value)
+    v = _builtin_function_to_json(value)
   elif inspect.isfunction(value):
-    return _function_to_json(value)
+    v = _function_to_json(value)
   elif inspect.ismethod(value):
-    return _method_to_json(value)
+    v = _method_to_json(value)
   # pytype: disable=module-attr
   elif isinstance(value, typing._Final):  # pylint: disable=protected-access
     # pytype: enable=module-attr
-    return _annotation_to_json(value)
+    v = _annotation_to_json(value)
   elif value is ...:
-    return {JSONConvertible.TYPE_NAME_KEY: 'type', 'name': 'builtins.Ellipsis'}
+    v = {JSONConvertible.TYPE_NAME_KEY: 'type', 'name': 'builtins.Ellipsis'}
   else:
+    v, converted = None, False
     if JSONConvertible.TYPE_CONVERTER is not None:
       converter = JSONConvertible.TYPE_CONVERTER(type(value))   # pylint: disable=not-callable
       if converter:
-        return to_json(converter(value))
-    return _OpaqueObject(value).to_json(**kwargs)
-
-
-def from_json(json_value: JSONValueType,
-              *,
-              force_dict: bool = False,
-              **kwargs) -> Any:
+        v = to_json(converter(value))
+        converted = True
+    if not converted:
+      v = _OpaqueObject(value).to_json(**kwargs)
+  return v
+
+
+def from_json(
+    json_value: JSONValueType,
+    *,
+    auto_import: bool = True,
+    auto_dict: bool = False,
+    **kwargs) -> Any:
   """Deserializes a (maybe) JSONConvertible value from JSON value.
 
   Args:
     json_value: Input JSON value.
-    force_dict: If True, "_type" keys will be stripped before loading. As a
-      result, JSONConvertible objects will be returned as dict.
+    auto_import: If True, when a '_type' is not registered, PyGlove will
+      identify its parent module and automatically import it. For example,
+      if the type is 'foo.bar.A', PyGlove will try to import 'foo.bar' and
+      find the class 'A' within the imported module.
+    auto_dict: If True, dict with '_type' that cannot be loaded will remain
+      as dict, with '_type' renamed to 'type_name'.
     **kwargs: Keyword arguments that will be passed to JSONConvertible.__init__.
 
   Returns:
     Deserialized value.
   """
-  if force_dict:
-    json_value = strip_types(json_value)
+  typename_resolved = kwargs.pop('_typename_resolved', False)
+  if not typename_resolved:
+    json_value = resolve_typenames(json_value, auto_import, auto_dict)
+
+  def child_from(v):
+    return from_json(v, _typename_resolved=True, **kwargs)
 
   if isinstance(json_value, list):
     if json_value and json_value[0] == JSONConvertible.TUPLE_MARKER:
@@ -411,41 +479,89 @@ def from_json(json_value: JSONValueType,
             f'Tuple should have at least one element '
             f'besides \'{JSONConvertible.TUPLE_MARKER}\'. '
             f'Encountered: {json_value}.')
-      return tuple([from_json(v, **kwargs) for v in json_value[1:]])
-    return [from_json(v, **kwargs) for v in json_value]
+      return tuple([child_from(v) for v in json_value[1:]])
+    return [child_from(v) for v in json_value]
   elif isinstance(json_value, dict):
     if JSONConvertible.TYPE_NAME_KEY not in json_value:
-      return {k: from_json(v, **kwargs) for k, v in json_value.items()}
-    type_name = json_value.pop(JSONConvertible.TYPE_NAME_KEY)
+      return {k: child_from(v) for k, v in json_value.items()}
+    factory_fn = json_value.pop(JSONConvertible.TYPE_NAME_KEY)
+    assert factory_fn is not None
+    return factory_fn(json_value, **kwargs)
+  return json_value
+
+
+def resolve_typenames(
+    json_value: JSONValueType,
+    auto_import: bool = True,
+    auto_dict: bool = False
+) -> JSONValueType:
+  """Inplace resolves the "_type" keys with their factories in a JSON tree."""
+
+  def _resolve_typename(v: Dict[str, Any]) -> bool:
+    """Returns True if the subtree is resolved for the first time."""
+    if JSONConvertible.TYPE_NAME_KEY not in v:
+      return True
+    if not isinstance(v[JSONConvertible.TYPE_NAME_KEY], str):
+      return False
+    type_name = v[JSONConvertible.TYPE_NAME_KEY]
     if type_name == 'type':
-      return _type_from_json(json_value)
+      factory_fn = _type_from_json
     elif type_name == 'function':
-      return _function_from_json(json_value)
+      factory_fn = _function_from_json
     elif type_name == 'method':
-      return _method_from_json(json_value)
+      factory_fn = _method_from_json
     else:
       cls = JSONConvertible.class_from_typename(type_name)
       if cls is None:
+        if auto_import:
+          try:
+            cls = _load_symbol(type_name)
+            assert inspect.isclass(cls), cls
+          except (ModuleNotFoundError, AttributeError) as e:
+            if not auto_dict:
+              raise TypeError(
+                  f'Cannot load class {type_name!r}.\n'
+                  'Try pass `auto_dict=True` to load the object into a dict '
+                  'without depending on the type.'
+              ) from e
+        elif not auto_dict:
+          raise TypeError(
+              f'Type name \'{type_name}\' is not registered '
+              'with a `pg.JSONConvertible` subclass.\n'
+              'Try pass `auto_import=True` to load the type from its module, '
+              'or pass `auto_dict=True` to load the object into a dict '
+              'without depending on the type.'
+          )
+
+      factory_fn = getattr(cls, 'from_json', None)
+      if cls is not None and factory_fn is None and not auto_dict:
         raise TypeError(
-            f'Type name \'{type_name}\' is not registered '
-            f'with a `pg.JSONConvertible` subclass.')
-      return cls.from_json(json_value, **kwargs)
-  return json_value
+            f'{cls} is not a `pg.JSONConvertible` subclass.'
+            'Try pass `auto_dict=True` to load the object into a dict '
+            'without depending on the type.'
+        )
+
+      if factory_fn is None and auto_dict:
+        v['type_name'] = type_name
+        v.pop(JSONConvertible.TYPE_NAME_KEY)
+        return True
+      assert factory_fn is not None
 
+    v[JSONConvertible.TYPE_NAME_KEY] = factory_fn
+    return True
 
-def strip_types(json_value: JSONValueType) -> JSONValueType:
-  """Inplace strips the "_type" key from a JSON tree."""
-  def _strip_type(v) -> None:
+  def _visit(v) -> None:
     if isinstance(v, (tuple, list)):
       for x in v:
-        _strip_type(x)
+        _visit(x)
     elif isinstance(v, dict):
-      if JSONConvertible.TYPE_NAME_KEY in v:
-        v.pop(JSONConvertible.TYPE_NAME_KEY)
-      for x in v.values():
-        _strip_type(x)
+      if _resolve_typename(v):
+        # Only resolve children when _types in this tree is not resolved
+        # previously
+        for x in v.values():
+          _visit(x)
 
-  _strip_type(json_value)
+  _visit(json_value)
   return json_value
 
 
@@ -610,6 +726,9 @@ def _load_symbol(type_name: str) -> Any:
     module_name = '.'.join(maybe_modules[:module_end_pos])
     parent_symbols = maybe_modules[module_end_pos:]
 
+  if not module_name:
+    raise ModuleNotFoundError(f'Cannot load symbol {type_name!r}.')
+
   # Import module and lookup parent symbols.
   module = importlib.import_module(module_name)
   parent = module
@@ -622,21 +741,26 @@ def _load_symbol(type_name: str) -> Any:
   return symbol
 
 
-def _type_from_json(json_value: Dict[str, str]) -> Type[Any]:
+def _type_from_json(json_value: Dict[str, str], **kwargs) -> Type[Any]:
   """Loads a type from a JSON dict."""
+  del kwargs
   t = _load_symbol(json_value['name'])
   if 'args' in json_value:
-    return _bind_type_args(t, from_json(json_value['args']))
+    return _bind_type_args(
+        t, from_json(json_value['args'], _typename_resolved=True)
+    )
   return t
 
 
-def _function_from_json(json_value: Dict[str, str]) -> types.FunctionType:
+def _function_from_json(
+    json_value: Dict[str, str], **kwargs) -> types.FunctionType:
   """Loads a function from a JSON dict."""
+  del kwargs
   function_name = json_value['name']
   if 'code' in json_value:
     code = marshal.loads(
         base64.decodebytes(json_value['code'].encode('utf-8')))
-    defaults = from_json(json_value['defaults'])
+    defaults = from_json(json_value['defaults'], _typename_resolved=True)
     return types.FunctionType(
         code=code,
         globals=globals(),
@@ -646,8 +770,9 @@ def _function_from_json(json_value: Dict[str, str]) -> types.FunctionType:
     return _load_symbol(function_name)
 
 
-def _method_from_json(json_value: Dict[str, str]) -> types.MethodType:
+def _method_from_json(json_value: Dict[str, str], **kwargs) -> types.MethodType:
   """Loads a class method from a JSON dict."""
+  del kwargs
   return _load_symbol(json_value['name'])
 
 

pyglove/core/{object_utils → utils}/json_conversion_test.py

@@ -11,13 +11,11 @@
 # 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.object_utils.json_conversion."""
-
 import abc
 import typing
 import unittest
-from pyglove.core.object_utils import json_conversion
 from pyglove.core.typing import inspect as pg_inspect
+from pyglove.core.utils import json_conversion
 
 
 class X:
@@ -96,7 +94,7 @@ class JSONConvertibleTest(unittest.TestCase):
         self.x = x
 
       def to_json(self):
-        return A.to_json_dict({
+        return self.__class__.to_json_dict({
             'x': self.x
         })
 
@@ -130,6 +128,27 @@ class JSONConvertibleTest(unittest.TestCase):
         (typename(B), C),
         list(json_conversion.JSONConvertible.registered_types()))
 
+    # Test load_types_for_deserialization.
+    class D(C):
+      auto_register = False
+
+    with self.assertRaisesRegex(TypeError, 'Cannot load class .*'):
+      json_conversion.from_json(D(1).to_json())
+
+    self.assertEqual(
+        json_conversion.JSONConvertible._TYPE_REGISTRY._ondemand_registry_stack,
+        []
+    )
+    with json_conversion.JSONConvertible.load_types_for_deserialization(A):
+      with json_conversion.JSONConvertible.load_types_for_deserialization(
+          D) as ondemand_registry:
+        self.assertEqual(ondemand_registry, {'A': A, 'D': D})
+        self.assertIsNotNone(json_conversion.from_json(D(1).to_json()))
+    self.assertEqual(
+        json_conversion.JSONConvertible._TYPE_REGISTRY._ondemand_registry_stack,
+        []
+    )
+
   def test_json_conversion(self):
 
     class T(json_conversion.JSONConvertible):
@@ -171,21 +190,41 @@ class JSONConvertibleTest(unittest.TestCase):
         T()
     )
 
-    # Test bad cases.
-    with self.assertRaisesRegex(
-        ValueError, 'Tuple should have at least one element besides .*'):
-      json_conversion.from_json(['__tuple__'])
-
-    with self.assertRaisesRegex(
-        TypeError, 'Type name .* is not registered'):
-      json_conversion.from_json({'_type': '__main__.ABC'})
-
   def assert_conversion_is(self, v):
     self.assertIs(json_conversion.from_json(json_conversion.to_json(v)), v)
 
   def assert_conversion_equal(self, v):
     self.assertEqual(json_conversion.from_json(json_conversion.to_json(v)), v)
 
+  class CustomJsonConvertible(json_conversion.JSONConvertible):
+    auto_register = False
+
+    def __init__(self, x=None):
+      self.x = x
+
+    def to_json(self):
+      return self.to_json_dict(
+          dict(x=(self.x, None)), exclude_default=True
+      )
+
+    def __eq__(self, other):
+      return isinstance(other, self.__class__) and self.x == other.x
+
+    def __ne__(self, other):
+      return not self.__eq__(other)
+
+  def test_json_conversion_with_auto_import(self):
+    json_dict = json_conversion.to_json(self.CustomJsonConvertible(1))
+
+    with self.assertRaisesRegex(
+        TypeError, 'Type name .* is not registered'):
+      json_conversion.from_json(json_dict, auto_import=False)
+
+    self.assertEqual(
+        json_conversion.from_json(json_dict),
+        self.CustomJsonConvertible(1)
+    )
+
   def test_json_conversion_for_types(self):
     # Built-in types.
     self.assert_conversion_is(int)
@@ -297,7 +336,8 @@ class JSONConvertibleTest(unittest.TestCase):
         ValueError, 'Cannot decode opaque object with pickle.'):
       json_conversion.from_json(json_dict)
 
-  def test_json_conversion_force_dict(self):
+  def test_json_conversion_auto_dict(self):
+    # Does not exist.
     self.assertEqual(
         json_conversion.from_json([
             '__tuple__',
@@ -306,12 +346,53 @@ class JSONConvertibleTest(unittest.TestCase):
                 '_type': 'Unknown type',
                 'x': [{
                     '_type': 'Unknown type',
+                }, {
+                    '_type': 'function',
+                    'name': 'builtins.print'
                 }]
             }
-        ], force_dict=True),
-        (1, {'x': [{}]})
+        ], auto_dict=True),
+        (1, {
+            'type_name': 'Unknown type',
+            'x': [{
+                'type_name': 'Unknown type',
+            }, print]
+        })
     )
 
+  def test_json_conversion_with_bad_types(self):
+    # Bad tuple.
+    with self.assertRaisesRegex(
+        ValueError, 'Tuple should have at least one element besides .*'):
+      json_conversion.from_json(['__tuple__'])
+
+    # Unregistered type without auto_import.
+    with self.assertRaisesRegex(
+        TypeError, 'Type name .* is not registered with'
+    ):
+      json_conversion.from_json(
+          {
+              '_type': 'Unknown type',
+              'x': [{
+                  '_type': 'Unknown type',
+              }]
+          }, auto_import=False
+      )
+
+    # Type does not exist.
+    with self.assertRaisesRegex(
+        TypeError, 'Cannot load class .*'):
+      json_conversion.from_json({'_type': '__main__.ABC'})
+
+    # Type exist but not a JSONConvertible subclass.
+    class A:
+      pass
+
+    json_conversion.JSONConvertible.register('__main__.A', A)
+    with self.assertRaisesRegex(
+        TypeError, '.* is not a `pg.JSONConvertible` subclass'):
+      json_conversion.from_json({'_type': '__main__.A'})
+
 
 if __name__ == '__main__':
   unittest.main()

pyglove/core/{object_utils → utils}/missing.py

@@ -14,11 +14,11 @@
 """Representing missing value for a field."""
 
 from typing import Any, Dict
-from pyglove.core.object_utils import common_traits
-from pyglove.core.object_utils import json_conversion
+from pyglove.core.utils import formatting
+from pyglove.core.utils import json_conversion
 
 
-class MissingValue(common_traits.Formattable, json_conversion.JSONConvertible):
+class MissingValue(formatting.Formattable, json_conversion.JSONConvertible):
   """Value placeholder for an unassigned attribute."""
 
   def format(self, *args, **kwargs):  # pytype: disable=signature-mismatch

pyglove/core/{object_utils → utils}/missing_test.py

@@ -11,11 +11,9 @@
 # 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.object_utils.missing."""
-
 import unittest
-from pyglove.core.object_utils import json_conversion
-from pyglove.core.object_utils import missing
+from pyglove.core.utils import json_conversion
+from pyglove.core.utils import missing
 
 
 class MissingValueTest(unittest.TestCase):

pyglove/core/utils/text_color.py

@@ -0,0 +1,128 @@
+# Copyright 2025 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Utility library for text coloring."""
+
+import re
+from typing import List, Optional
+
+try:
+  import termcolor      # pylint: disable=g-import-not-at-top
+except ImportError:
+  termcolor = None
+
+
+# Regular expression for ANSI color characters.
+_ANSI_COLOR_REGEX = re.compile(r'\x1b\[[0-9;]*m')
+
+
+def decolor(text: str) -> str:
+  """De-colors a string that may contains ANSI color characters."""
+  return re.sub(_ANSI_COLOR_REGEX, '', text)
+
+
+def colored(
+    text: str,
+    color: Optional[str] = None,
+    background: Optional[str] = None,
+    styles: Optional[List[str]] = None,
+) -> str:
+  """Returns the colored text with ANSI color characters.
+
+  Args:
+    text: A string that may or may not already has ANSI color characters.
+    color: A string for text colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    background: A string for background colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    styles: A list of strings for applying styles on the text.
+      Applicable values are:
+      'bold', 'dark', 'underline', 'blink', 'reverse', 'concealed'.
+
+  Returns:
+    A string with ANSI color characters embracing the entire text.
+  """
+  if not termcolor:
+    return text
+  return termcolor.colored(
+      text,
+      color=color,
+      on_color=('on_' + background) if background else None,
+      attrs=styles
+  )
+
+
+def colored_block(
+    text: str,
+    block_start: str,
+    block_end: str,
+    color: Optional[str] = None,
+    background: Optional[str] = None,
+    styles: Optional[List[str]] = None,
+) -> str:
+  """Apply colors to text blocks.
+
+  Args:
+    text: A string that may or may not already has ANSI color characters.
+    block_start: A string that signals the start of a block. E.g. '{{'
+    block_end: A string that signals the end of a block. E.g. '}}'.
+    color: A string for text colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    background: A string for background colors. Applicable values are:
+      'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'.
+    styles: A list of strings for applying styles on the text.
+      Applicable values are:
+      'bold', 'dark', 'underline', 'blink', 'reverse', 'concealed'.
+
+  Returns:
+    A string with ANSI color characters embracing the matched text blocks.
+  """
+  if not color and not background and not styles:
+    return text
+
+  s = []
+  start_index = 0
+  end_index = 0
+  previous_color = None
+
+  def write_nonblock_text(text: str, previous_color: Optional[str]):
+    if previous_color:
+      s.append(previous_color)
+    s.append(text)
+
+  while start_index < len(text):
+    start_index = text.find(block_start, end_index)
+    if start_index == -1:
+      write_nonblock_text(text[end_index:], previous_color)
+      break
+
+    # Deal with text since last block.
+    since_last_block = text[end_index:start_index]
+    write_nonblock_text(since_last_block, previous_color)
+    colors = re.findall(_ANSI_COLOR_REGEX, since_last_block)
+    if colors:
+      previous_color = colors[-1]
+
+    # Match block.
+    end_index = text.find(block_end, start_index + len(block_start))
+    if end_index == -1:
+      write_nonblock_text(text[start_index:], previous_color)
+      break
+    end_index += len(block_end)
+
+    # Write block text.
+    block = text[start_index:end_index]
+    block = colored(
+        block, color=color, background=background, styles=styles)
+    s.append(block)
+  return ''.join(s)