File: save_tree2file.py

package info (click to toggle)
treelib 1.8.0-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 488 kB
  • sloc: python: 3,679; makefile: 85; sh: 37
file content (233 lines) | stat: -rw-r--r-- 6,853 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
#!/usr/bin/env python
"""
Save Tree to File Example - Demonstrate tree export capabilities

This example shows how to:
- Save trees to text files
- Export trees in different formats
- Load and manipulate saved tree data
- Export to various file formats (JSON, Graphviz DOT)

Author: treelib contributors
"""

import os
import tempfile

from treelib import Tree


def create_sample_tree():
    """
    Create a sample tree for demonstration.

    Returns:
        Tree: A populated tree
    """
    tree = Tree()

    # Create a simple organizational structure
    tree.create_node("CEO", "ceo")
    tree.create_node("CTO", "cto", parent="ceo")
    tree.create_node("CFO", "cfo", parent="ceo")
    tree.create_node("VP Engineering", "vp_eng", parent="cto")
    tree.create_node("VP Finance", "vp_fin", parent="cfo")
    tree.create_node("Senior Developer", "dev1", parent="vp_eng")
    tree.create_node("Junior Developer", "dev2", parent="vp_eng")
    tree.create_node("Accountant", "acc1", parent="vp_fin")

    return tree


def demonstrate_text_export():
    """Demonstrate saving tree to text file."""
    print("=" * 60)
    print("๐Ÿ“„ SAVING TREE TO TEXT FILE")
    print("=" * 60)

    tree = create_sample_tree()

    print("๐ŸŒณ Original tree structure:")
    tree.show()

    # Save to text file
    filename = "organization_tree.txt"
    tree.save2file(filename)

    print(f"\n๐Ÿ’พ Tree saved to '{filename}'")

    # Display file contents
    print("\n๐Ÿ“– File contents:")
    try:
        with open(filename, "r", encoding="utf-8") as f:
            content = f.read()
            print(content)
    except FileNotFoundError:
        print(f"โŒ Error: Could not read file '{filename}'")

    # Clean up
    try:
        os.remove(filename)
        print(f"๐Ÿงน Cleaned up temporary file '{filename}'")
    except FileNotFoundError:
        pass


def demonstrate_json_export():
    """Demonstrate exporting tree to JSON format."""
    print("\n" + "=" * 60)
    print("๐Ÿ“„ EXPORTING TREE TO JSON")
    print("=" * 60)

    tree = create_sample_tree()

    # Export to JSON
    json_output = tree.to_json()
    print("๐Ÿ“‹ Tree as JSON:")
    print(json_output)

    # Save JSON to file
    json_filename = "organization_tree.json"
    with open(json_filename, "w", encoding="utf-8") as f:
        f.write(json_output)

    print(f"\n๐Ÿ’พ JSON saved to '{json_filename}'")

    # Clean up
    try:
        os.remove(json_filename)
        print(f"๐Ÿงน Cleaned up temporary file '{json_filename}'")
    except FileNotFoundError:
        pass


def demonstrate_dict_export():
    """Demonstrate converting tree to dictionary format."""
    print("\n" + "=" * 60)
    print("๐Ÿ“„ CONVERTING TREE TO DICTIONARY")
    print("=" * 60)

    tree = create_sample_tree()

    # Convert to dictionary
    tree_dict = tree.to_dict()
    print("๐Ÿ“‹ Tree as dictionary:")
    print(tree_dict)

    # Convert to dictionary with data
    print("\n๐Ÿ“‹ Tree as dictionary (with custom data):")
    # Add some custom data to nodes
    tree["ceo"].data = {"salary": 200000, "department": "Executive"}
    tree["cto"].data = {"salary": 150000, "department": "Technology"}
    tree["cfo"].data = {"salary": 140000, "department": "Finance"}

    tree_dict_with_data = tree.to_dict(with_data=True)
    print(tree_dict_with_data)


def demonstrate_custom_formatting():
    """Demonstrate custom tree formatting and export."""
    print("\n" + "=" * 60)
    print("๐Ÿ“„ CUSTOM TREE FORMATTING")
    print("=" * 60)

    tree = create_sample_tree()

    # Add custom data to demonstrate data property display
    class JobInfo:
        def __init__(self, title, level):
            self.title = title
            self.level = level

    tree["ceo"].data = JobInfo("Chief Executive Officer", 1)
    tree["cto"].data = JobInfo("Chief Technology Officer", 2)
    tree["cfo"].data = JobInfo("Chief Financial Officer", 2)
    tree["vp_eng"].data = JobInfo("Vice President of Engineering", 3)
    tree["vp_fin"].data = JobInfo("Vice President of Finance", 3)
    tree["dev1"].data = JobInfo("Senior Developer", 4)
    tree["dev2"].data = JobInfo("Junior Developer", 4)
    tree["acc1"].data = JobInfo("Accountant", 4)

    print("๐ŸŒณ Tree with custom formatting:")

    # Use show() with custom parameters
    output = tree.show(
        idhidden=False,  # Show node IDs
        line_type="ascii-em",  # Fancy line style
        data_property="title",  # Show custom data property
        stdout=False,  # Return as string instead of printing
    )
    print(output)

    # Save formatted output to file
    formatted_filename = "formatted_tree.txt"
    tree.save2file(formatted_filename, idhidden=False, line_type="ascii-em", data_property="title")

    print(f"\n๐Ÿ’พ Formatted tree saved to '{formatted_filename}'")

    # Clean up
    try:
        os.remove(formatted_filename)
        print(f"๐Ÿงน Cleaned up temporary file '{formatted_filename}'")
    except FileNotFoundError:
        pass


def demonstrate_graphviz_export():
    """Demonstrate exporting tree to Graphviz DOT format."""
    print("\n" + "=" * 60)
    print("๐Ÿ“„ EXPORTING TO GRAPHVIZ DOT FORMAT")
    print("=" * 60)

    tree = create_sample_tree()

    # Create temporary file for DOT output
    with tempfile.NamedTemporaryFile(mode="w", suffix=".dot", delete=False) as tmp_file:
        dot_filename = tmp_file.name

    print("๐ŸŒณ Exporting tree to Graphviz DOT format...")
    tree.to_graphviz(filename=dot_filename, shape="box")

    print(f"๐Ÿ’พ DOT file saved to '{dot_filename}'")

    # Display DOT file contents
    print("\n๐Ÿ“– DOT file contents:")
    try:
        with open(dot_filename, "r", encoding="utf-8") as f:
            content = f.read()
            print(content)
    except FileNotFoundError:
        print(f"โŒ Error: Could not read DOT file '{dot_filename}'")

    print("\n๐Ÿ’ก Tip: You can use Graphviz tools to convert this to SVG, PNG, or PDF:")
    print(f"    dot -Tsvg {dot_filename} -o tree.svg")
    print(f"    dot -Tpng {dot_filename} -o tree.png")

    # Clean up
    try:
        os.remove(dot_filename)
        print(f"\n๐Ÿงน Cleaned up temporary file '{dot_filename}'")
    except FileNotFoundError:
        pass


def main():
    """Main function demonstrating tree export capabilities."""
    print("๐Ÿ’พ Welcome to the TreeLib Tree Export Example!")
    print("This example shows different ways to save and export trees.")

    # Run all demonstrations
    demonstrate_text_export()
    demonstrate_json_export()
    demonstrate_dict_export()
    demonstrate_custom_formatting()
    demonstrate_graphviz_export()

    print("\n" + "=" * 60)
    print("๐ŸŽ‰ All export examples completed!")
    print("๐Ÿ’ก Try experimenting with different formatting options and export methods.")
    print("=" * 60)


if __name__ == "__main__":
    main()