Add indent contextmanager

Author: spenczar

protogen/__init__.py

@@ -26,7 +26,7 @@
     def generate(gen: protogen.Plugin):
         for f in gen.files_to_generate:
             g = gen.new_generated_file(
-                f.proto.name.replace(".proto", ".py"), 
+                f.proto.name.replace(".proto", ".py"),
                 f.py_import_path,
             )
             g.P("# Generated code ahead.")
@@ -49,10 +49,11 @@ def generate(gen: protogen.Plugin):
 
 """
 
+import contextlib
 import enum
 import keyword
 import sys
-from typing import BinaryIO, Callable, Dict, List, Optional, Set, Tuple
+from typing import BinaryIO, Callable, Dict, List, Iterator, Optional, Set, Tuple
 from operator import ior
 from functools import reduce
 
@@ -1483,6 +1484,34 @@ def set_indent(self, level: int) -> int:
         self._indent = level
         return old
 
+    @contextlib.contextmanager
+    def indent(self, level: int = 4) -> Iterator[None]:
+        """Increments the indentation within a contextmanager block.
+
+        Increases the indentation level for all calls to :func:`P` within a
+        contextmanager block. When the enclosed block exits, the indentation
+        level will be restored to its prior level.
+
+        Arguments
+        ---------
+        level : int
+            The amount to increment the indentation level by.
+
+        Example
+        -------
+        >>> g.P("class Myclass:")
+        >>> with g.indent():
+        ...     g.P("def __init__():")
+        ...     with g.indent():
+        ...         g.P("pass")
+        """
+        old = self._indent
+        self._indent += level
+        try:
+            yield
+        finally:
+            self._indent = old
+
     def P(self, *args):
         """Add a new line to the output buffer.
 

test/plugin/main.py

@@ -65,13 +65,11 @@ def generate_message(g: protogen.GeneratedFile, message: protogen.Message, lvl:
 
 
 def generate_enum(g: protogen.GeneratedFile, e: protogen.Enum, lvl: int):
-    reset = g.set_indent(lvl)
-    g.P("- name: ", e.full_name)
-    g.P("  values:")
-    for v in e.values:
-        g.P("  - name: ", v.full_name)
-
-    g.set_indent(reset)
+    with g.indent(lvl):
+        g.P("- name: ", e.full_name)
+        g.P("  values:")
+        for v in e.values:
+            g.P("  - name: ", v.full_name)
 
 
 def generate_service(g: protogen.Service, s: protogen.Service):
@@ -94,7 +92,7 @@ def generate_extensions(g: protogen.GeneratedFile, e: protogen.Extension):
 
 
 def collect_messages(
-    fm: Union[protogen.File, protogen.Message]
+    fm: Union[protogen.File, protogen.Message],
 ) -> List[protogen.Message]:
     messages = []
     for m in fm.messages:

test/test_protogen.py

@@ -1,3 +1,4 @@
+import protogen
 from protogen.test import run_plugin
 
 
@@ -43,3 +44,23 @@ def test_parameter():
         golden = f.read()
 
     assert golden == generated
+
+
+def test_with_indent():
+    g = protogen.GeneratedFile("test_with_indent", "")
+    g.P("top-level")
+    with g.indent(2):
+        g.P("indented-by-two")
+        with g.indent(4):
+            g.P("indented-by-six")
+        g.P("return-by-two")
+    g.P("return-top-level")
+
+    expected = [
+        "top-level",
+        "  indented-by-two",
+        "      indented-by-six",
+        "  return-by-two",
+        "return-top-level",
+    ]
+    assert g._buf == expected