Add indent control knob (#49)

Co-authored-by: Taneli Hukkinen <[email protected]>

Author: wimglenn

src/tomli_w/_writer.py

@@ -11,7 +11,6 @@
 ILLEGAL_BASIC_STR_CHARS = frozenset('"\\') | ASCII_CTRL - frozenset("\t")
 BARE_KEY_CHARS = frozenset(string.ascii_letters + string.digits + "-_")
 ARRAY_TYPES = (list, tuple)
-ARRAY_INDENT = " " * 4
 MAX_LINE_LENGTH = 100
 
 COMPACT_ESCAPES = MappingProxyType(
@@ -27,22 +26,29 @@
 
 
 def dump(
-    __obj: Mapping[str, Any], __fp: IO[bytes], *, multiline_strings: bool = False
+    __obj: Mapping[str, Any],
+    __fp: IO[bytes],
+    *,
+    multiline_strings: bool = False,
+    indent: int = 4,
 ) -> None:
-    ctx = Context(multiline_strings, {})
+    ctx = Context(multiline_strings, {}, " " * indent)
     for chunk in gen_table_chunks(__obj, ctx, name=""):
         __fp.write(chunk.encode())
 
 
-def dumps(__obj: Mapping[str, Any], *, multiline_strings: bool = False) -> str:
-    ctx = Context(multiline_strings, {})
+def dumps(
+    __obj: Mapping[str, Any], *, multiline_strings: bool = False, indent: int = 4
+) -> str:
+    ctx = Context(multiline_strings, {}, " " * indent)
     return "".join(gen_table_chunks(__obj, ctx, name=""))
 
 
 class Context(NamedTuple):
     allow_multiline: bool
     # cache rendered inline tables (mapping from object id to rendered inline table)
     inline_table_cache: dict[int, str]
+    indent_str: str
 
 
 def gen_table_chunks(
@@ -136,8 +142,8 @@ def format_inline_table(obj: Mapping, ctx: Context) -> str:
 def format_inline_array(obj: tuple | list, ctx: Context, nest_level: int) -> str:
     if not obj:
         return "[]"
-    item_indent = ARRAY_INDENT * (1 + nest_level)
-    closing_bracket_indent = ARRAY_INDENT * nest_level
+    item_indent = ctx.indent_str * (1 + nest_level)
+    closing_bracket_indent = ctx.indent_str * nest_level
     return (
         "[\n"
         + ",\n".join(
@@ -197,5 +203,5 @@ def is_aot(obj: Any) -> bool:
 def is_suitable_inline_table(obj: Mapping, ctx: Context) -> bool:
     """Use heuristics to decide if the inline-style representation is a good
     choice for a given table."""
-    rendered_inline = f"{ARRAY_INDENT}{format_inline_table(obj, ctx)},"
+    rendered_inline = f"{ctx.indent_str}{format_inline_table(obj, ctx)},"
     return len(rendered_inline) <= MAX_LINE_LENGTH and "\n" not in rendered_inline

tests/test_style.py

@@ -272,3 +272,16 @@ def test_multiline_in_aot():
 ]
 """
     )
+
+
+def test_array_indent_override():
+    data = {"k0": ["v1", "v2"]}
+    assert (
+        tomli_w.dumps(data, indent=2)
+        == """\
+k0 = [
+  "v1",
+  "v2",
+]
+"""
+    )