Skip to content

Commit 61e6e54

Browse files
authored
Merge pull request #286 from mmcdermott/285_print_to_file
Added print kwargs to print options.
2 parents 2eeb9c8 + 2894b08 commit 61e6e54

File tree

1 file changed

+32
-2
lines changed

1 file changed

+32
-2
lines changed

bigtree/tree/export.py

+32-2
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ def print_tree(
7373
attr_omit_null: bool = False,
7474
attr_bracket: List[str] = ["[", "]"],
7575
style: Union[str, Iterable[str], BasePrintStyle] = "const",
76+
**print_kwargs,
7677
) -> None:
7778
"""Print tree to console, starting from `tree`.
7879
@@ -90,6 +91,8 @@ def print_tree(
9091
- (BasePrintStyle): `ANSIPrintStyle`, `ASCIIPrintStyle`, `ConstPrintStyle`, `ConstBoldPrintStyle`, `RoundedPrintStyle`,
9192
`DoublePrintStyle` style or inherit from `BasePrintStyle`
9293
94+
Remaining kwargs are passed without modification to python's `print` function.
95+
9396
Examples:
9497
**Printing tree**
9598
@@ -188,6 +191,19 @@ def print_tree(
188191
| `-- e
189192
`-- c
190193
194+
**Printing to a file**
195+
196+
>>> import io
197+
>>> output = io.StringIO()
198+
>>> print_tree(root, file=output)
199+
>>> string = output.getvalue()
200+
>>> print(string)
201+
a
202+
├── b
203+
│ ├── d
204+
│ └── e
205+
└── c
206+
191207
Args:
192208
tree (Node): tree to print
193209
node_name_or_path (str): node to print from, becomes the root node of printing
@@ -232,7 +248,7 @@ def print_tree(
232248
if attr_str:
233249
attr_str = f" {attr_bracket_open}{attr_str}{attr_bracket_close}"
234250
node_str = f"{_node.node_name}{attr_str}"
235-
print(f"{pre_str}{fill_str}{node_str}")
251+
print(f"{pre_str}{fill_str}{node_str}", **print_kwargs)
236252

237253

238254
def yield_tree(
@@ -416,6 +432,7 @@ def hprint_tree(
416432
max_depth: int = 0,
417433
intermediate_node_name: bool = True,
418434
style: Union[str, Iterable[str], BaseHPrintStyle] = "const",
435+
**print_kwargs,
419436
) -> None:
420437
"""Print tree in horizontal orientation to console, starting from `tree`.
421438
@@ -430,6 +447,8 @@ def hprint_tree(
430447
- (BaseHPrintStyle): `ANSIHPrintStyle`, `ASCIIHPrintStyle`, `ConstHPrintStyle`, `ConstBoldHPrintStyle`,
431448
`RoundedHPrintStyle`, `DoubleHPrintStyle` style or inherit from BaseHPrintStyle
432449
450+
Remaining kwargs are passed without modification to python's `print` function.
451+
433452
Examples:
434453
**Printing tree**
435454
@@ -504,6 +523,17 @@ def hprint_tree(
504523
- a -+ \\- e
505524
\\- c
506525
526+
**Printing to a file**
527+
>>> import io
528+
>>> output = io.StringIO()
529+
>>> hprint_tree(root, file=output)
530+
>>> string = output.getvalue()
531+
>>> print(string)
532+
┌─ d
533+
┌─ b ─┤
534+
─ a ─┤ └─ e
535+
└─ c
536+
507537
Args:
508538
tree (Node): tree to print
509539
node_name_or_path (str): node to print from, becomes the root node of printing
@@ -518,7 +548,7 @@ def hprint_tree(
518548
max_depth=max_depth,
519549
style=style,
520550
)
521-
print("\n".join(result))
551+
print("\n".join(result), **print_kwargs)
522552

523553

524554
def hyield_tree(

0 commit comments

Comments
 (0)