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", +] +""" + )