From af607ff46c1226f0bb443f2ac14b5575bd2d4117 Mon Sep 17 00:00:00 2001
From: Wim Jeantine-Glenn <wim.glenn@gmail.com>
Date: Tue, 8 Oct 2024 03:08:03 -0500
Subject: [PATCH] Add `indent` control knob (#49)

Co-authored-by: Taneli Hukkinen <3275109+hukkin@users.noreply.github.com>
---
 src/tomli_w/_writer.py | 22 ++++++++++++++--------
 tests/test_style.py    | 13 +++++++++++++
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/src/tomli_w/_writer.py b/src/tomli_w/_writer.py
index e337dd5..84d2f1c 100644
--- a/src/tomli_w/_writer.py
+++ b/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,15 +26,21 @@
 
 
 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=""))
 
 
@@ -43,6 +48,7 @@ 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
diff --git a/tests/test_style.py b/tests/test_style.py
index 206f8a4..9cfc926 100644
--- a/tests/test_style.py
+++ b/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",
+]
+"""
+    )