diff --git a/protogen/__init__.py b/protogen/__init__.py index 72365c3..f43438d 100644 --- a/protogen/__init__.py +++ b/protogen/__init__.py @@ -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. diff --git a/test/plugin/main.py b/test/plugin/main.py index 4acdcab..a34866e 100755 --- a/test/plugin/main.py +++ b/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): diff --git a/test/test_protogen.py b/test/test_protogen.py index 695a929..7a5bfc4 100644 --- a/test/test_protogen.py +++ b/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