Skip to content

Commit

Permalink
Add indent control knob (#49)
Browse files Browse the repository at this point in the history
Co-authored-by: Taneli Hukkinen <3275109+hukkin@users.noreply.github.com>
  • Loading branch information
wimglenn and hukkin authored Oct 8, 2024
1 parent 36354bd commit af607ff
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 8 deletions.
22 changes: 14 additions & 8 deletions src/tomli_w/_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
13 changes: 13 additions & 0 deletions tests/test_style.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]
"""
)

0 comments on commit af607ff

Please sign in to comment.