1+ import io
2+
13import pandas as pd
24import polars as pl
35import pydot
2325from tests .node .test_node import assert_tree_structure_node_root
2426from tests .test_constants import Constants
2527
26- tree_node_str = """a [age=90]\n ├── b [age=65]\n │ ├── d [age=40]\n │ └── e [age=35]\n │ ├── g [age=10]
27- │ └── h [age=6]\n └── c [age=60]\n └── f [age=38]\n """
28- tree_node_no_attr_str = """a\n ├── b\n │ ├── d\n │ └── e\n │ ├── g\n │ └── h\n └── c\n └── f\n """
29-
28+ tree_node_str = (
29+ "a [age=90]\n "
30+ "├── b [age=65]\n "
31+ "│ ├── d [age=40]\n "
32+ "│ └── e [age=35]\n "
33+ "│ ├── g [age=10]\n "
34+ "│ └── h [age=6]\n "
35+ "└── c [age=60]\n "
36+ " └── f [age=38]\n "
37+ )
38+ tree_node_no_attr_str = (
39+ "a\n "
40+ "├── b\n "
41+ "│ ├── d\n "
42+ "│ └── e\n "
43+ "│ ├── g\n "
44+ "│ └── h\n "
45+ "└── c\n "
46+ " └── f\n "
47+ )
48+ tree_node_hstr = (
49+ " ┌─ d\n "
50+ " ┌─ b ─┤ ┌─ g\n "
51+ "─ a ─┤ └─ e ─┤\n "
52+ " │ └─ h\n "
53+ " └─ c ─── f\n "
54+ )
55+ # fmt: off
56+ tree_node_branch_hstr = (
57+ " ┌─ d\n "
58+ "─ b ─┤ ┌─ g\n "
59+ " └─ e ─┤\n "
60+ " └─ h\n "
61+ )
62+ # fmt: on
3063LOCAL = Constants .LOCAL
3164
3265
3366class TestPrintTree :
3467 @staticmethod
3568 def test_print_tree_child_node_name (tree_node ):
36- expected_str = "b\n " "├── d\n " "└── e\n " " ├── g\n " " └── h\n "
69+ # fmt: off
70+ expected_str = (
71+ "b\n "
72+ "├── d\n "
73+ "└── e\n "
74+ " ├── g\n "
75+ " └── h\n "
76+ )
77+ # fmt: on
3778 assert_print_statement (
3879 print_tree ,
3980 expected_str ,
@@ -54,7 +95,15 @@ def test_print_tree_child_node_name_error(tree_node):
5495
5596 @staticmethod
5697 def test_print_tree_child_node_path (tree_node ):
57- expected_str = "b\n " "├── d\n " "└── e\n " " ├── g\n " " └── h\n "
98+ # fmt: off
99+ expected_str = (
100+ "b\n "
101+ "├── d\n "
102+ "└── e\n "
103+ " ├── g\n "
104+ " └── h\n "
105+ )
106+ # fmt: on
58107 assert_print_statement (
59108 print_tree ,
60109 expected_str ,
@@ -352,20 +401,19 @@ def test_print_tree_custom_style_missing_style_error(tree_node):
352401 )
353402 assert str (exc_info .value ) == Constants .ERROR_NODE_EXPORT_PRINT_STYLE_SELECT
354403
404+ @staticmethod
405+ def test_print_tree_print_kwargs (tree_node ):
406+ output = io .StringIO ()
407+ print_tree (tree_node , file = output )
408+ assert output .getvalue () == tree_node_no_attr_str
409+
355410
356411class TestHPrintTree :
357412 @staticmethod
358413 def test_hprint_tree (tree_node ):
359- expected_str = (
360- " ┌─ d\n "
361- " ┌─ b ─┤ ┌─ g\n "
362- "─ a ─┤ └─ e ─┤\n "
363- " │ └─ h\n "
364- " └─ c ─── f\n "
365- )
366414 assert_print_statement (
367415 hprint_tree ,
368- expected_str ,
416+ tree_node_hstr ,
369417 tree = tree_node ,
370418 )
371419
@@ -473,12 +521,9 @@ def test_hprint_tree_diff_node_name_length(tree_node):
473521
474522 @staticmethod
475523 def test_hprint_tree_child_node_name (tree_node ):
476- expected_str = (
477- " ┌─ d\n " "─ b ─┤ ┌─ g\n " " └─ e ─┤\n " " └─ h\n "
478- )
479524 assert_print_statement (
480525 hprint_tree ,
481- expected_str ,
526+ tree_node_branch_hstr ,
482527 tree = tree_node ,
483528 node_name_or_path = "b" ,
484529 )
@@ -496,12 +541,9 @@ def test_hprint_tree_child_node_name_error(tree_node):
496541
497542 @staticmethod
498543 def test_hprint_tree_child_node_path (tree_node ):
499- expected_str = (
500- " ┌─ d\n " "─ b ─┤ ┌─ g\n " " └─ e ─┤\n " " └─ h\n "
501- )
502544 assert_print_statement (
503545 hprint_tree ,
504- expected_str ,
546+ tree_node_branch_hstr ,
505547 tree = tree_node ,
506548 node_name_or_path = "a/b" ,
507549 )
@@ -749,6 +791,12 @@ def test_hprint_tree_custom_style_missing_style_error(tree_node):
749791 hprint_tree (tree_node , style = ["-" , "+" , "+" , "|" , "-" , "=" ])
750792 assert str (exc_info .value ) == Constants .ERROR_NODE_EXPORT_HPRINT_STYLE_SELECT
751793
794+ @staticmethod
795+ def test_hprint_tree_print_kwargs (tree_node ):
796+ output = io .StringIO ()
797+ hprint_tree (tree_node , file = output )
798+ assert output .getvalue () == tree_node_hstr
799+
752800
753801class TestTreeToDataFrame :
754802 @staticmethod
0 commit comments