edsl 0.1.39.dev1__py3-none-any.whl → 0.1.39.dev3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (256) hide show
  1. edsl/Base.py +332 -332
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +49 -49
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +867 -867
  7. edsl/agents/AgentList.py +413 -413
  8. edsl/agents/Invigilator.py +233 -233
  9. edsl/agents/InvigilatorBase.py +270 -265
  10. edsl/agents/PromptConstructor.py +354 -354
  11. edsl/agents/__init__.py +3 -3
  12. edsl/agents/descriptors.py +99 -99
  13. edsl/agents/prompt_helpers.py +129 -129
  14. edsl/auto/AutoStudy.py +117 -117
  15. edsl/auto/StageBase.py +230 -230
  16. edsl/auto/StageGenerateSurvey.py +178 -178
  17. edsl/auto/StageLabelQuestions.py +125 -125
  18. edsl/auto/StagePersona.py +61 -61
  19. edsl/auto/StagePersonaDimensionValueRanges.py +88 -88
  20. edsl/auto/StagePersonaDimensionValues.py +74 -74
  21. edsl/auto/StagePersonaDimensions.py +69 -69
  22. edsl/auto/StageQuestions.py +73 -73
  23. edsl/auto/SurveyCreatorPipeline.py +21 -21
  24. edsl/auto/utilities.py +224 -224
  25. edsl/base/Base.py +279 -279
  26. edsl/config.py +157 -157
  27. edsl/conversation/Conversation.py +290 -290
  28. edsl/conversation/car_buying.py +58 -58
  29. edsl/conversation/chips.py +95 -95
  30. edsl/conversation/mug_negotiation.py +81 -81
  31. edsl/conversation/next_speaker_utilities.py +93 -93
  32. edsl/coop/PriceFetcher.py +54 -54
  33. edsl/coop/__init__.py +2 -2
  34. edsl/coop/coop.py +1028 -1028
  35. edsl/coop/utils.py +131 -131
  36. edsl/data/Cache.py +555 -555
  37. edsl/data/CacheEntry.py +233 -233
  38. edsl/data/CacheHandler.py +149 -149
  39. edsl/data/RemoteCacheSync.py +78 -78
  40. edsl/data/SQLiteDict.py +292 -292
  41. edsl/data/__init__.py +4 -4
  42. edsl/data/orm.py +10 -10
  43. edsl/data_transfer_models.py +73 -73
  44. edsl/enums.py +175 -175
  45. edsl/exceptions/BaseException.py +21 -21
  46. edsl/exceptions/__init__.py +54 -54
  47. edsl/exceptions/agents.py +42 -42
  48. edsl/exceptions/cache.py +5 -5
  49. edsl/exceptions/configuration.py +16 -16
  50. edsl/exceptions/coop.py +10 -10
  51. edsl/exceptions/data.py +14 -14
  52. edsl/exceptions/general.py +34 -34
  53. edsl/exceptions/jobs.py +33 -33
  54. edsl/exceptions/language_models.py +63 -63
  55. edsl/exceptions/prompts.py +15 -15
  56. edsl/exceptions/questions.py +91 -91
  57. edsl/exceptions/results.py +29 -29
  58. edsl/exceptions/scenarios.py +22 -22
  59. edsl/exceptions/surveys.py +37 -37
  60. edsl/inference_services/AnthropicService.py +87 -87
  61. edsl/inference_services/AwsBedrock.py +120 -120
  62. edsl/inference_services/AzureAI.py +217 -217
  63. edsl/inference_services/DeepInfraService.py +18 -18
  64. edsl/inference_services/GoogleService.py +148 -148
  65. edsl/inference_services/GroqService.py +20 -20
  66. edsl/inference_services/InferenceServiceABC.py +147 -147
  67. edsl/inference_services/InferenceServicesCollection.py +97 -97
  68. edsl/inference_services/MistralAIService.py +123 -123
  69. edsl/inference_services/OllamaService.py +18 -18
  70. edsl/inference_services/OpenAIService.py +224 -224
  71. edsl/inference_services/PerplexityService.py +163 -163
  72. edsl/inference_services/TestService.py +89 -89
  73. edsl/inference_services/TogetherAIService.py +170 -170
  74. edsl/inference_services/models_available_cache.py +118 -118
  75. edsl/inference_services/rate_limits_cache.py +25 -25
  76. edsl/inference_services/registry.py +41 -41
  77. edsl/inference_services/write_available.py +10 -10
  78. edsl/jobs/Answers.py +56 -56
  79. edsl/jobs/Jobs.py +898 -898
  80. edsl/jobs/JobsChecks.py +147 -147
  81. edsl/jobs/JobsPrompts.py +268 -268
  82. edsl/jobs/JobsRemoteInferenceHandler.py +239 -239
  83. edsl/jobs/__init__.py +1 -1
  84. edsl/jobs/buckets/BucketCollection.py +63 -63
  85. edsl/jobs/buckets/ModelBuckets.py +65 -65
  86. edsl/jobs/buckets/TokenBucket.py +251 -251
  87. edsl/jobs/interviews/Interview.py +661 -661
  88. edsl/jobs/interviews/InterviewExceptionCollection.py +99 -99
  89. edsl/jobs/interviews/InterviewExceptionEntry.py +186 -186
  90. edsl/jobs/interviews/InterviewStatistic.py +63 -63
  91. edsl/jobs/interviews/InterviewStatisticsCollection.py +25 -25
  92. edsl/jobs/interviews/InterviewStatusDictionary.py +78 -78
  93. edsl/jobs/interviews/InterviewStatusLog.py +92 -92
  94. edsl/jobs/interviews/ReportErrors.py +66 -66
  95. edsl/jobs/interviews/interview_status_enum.py +9 -9
  96. edsl/jobs/runners/JobsRunnerAsyncio.py +466 -466
  97. edsl/jobs/runners/JobsRunnerStatus.py +330 -330
  98. edsl/jobs/tasks/QuestionTaskCreator.py +242 -242
  99. edsl/jobs/tasks/TaskCreators.py +64 -64
  100. edsl/jobs/tasks/TaskHistory.py +450 -450
  101. edsl/jobs/tasks/TaskStatusLog.py +23 -23
  102. edsl/jobs/tasks/task_status_enum.py +163 -163
  103. edsl/jobs/tokens/InterviewTokenUsage.py +27 -27
  104. edsl/jobs/tokens/TokenUsage.py +34 -34
  105. edsl/language_models/KeyLookup.py +30 -30
  106. edsl/language_models/LanguageModel.py +668 -668
  107. edsl/language_models/ModelList.py +155 -155
  108. edsl/language_models/RegisterLanguageModelsMeta.py +184 -184
  109. edsl/language_models/__init__.py +3 -3
  110. edsl/language_models/fake_openai_call.py +15 -15
  111. edsl/language_models/fake_openai_service.py +61 -61
  112. edsl/language_models/registry.py +190 -190
  113. edsl/language_models/repair.py +156 -156
  114. edsl/language_models/unused/ReplicateBase.py +83 -83
  115. edsl/language_models/utilities.py +64 -64
  116. edsl/notebooks/Notebook.py +258 -258
  117. edsl/notebooks/__init__.py +1 -1
  118. edsl/prompts/Prompt.py +362 -362
  119. edsl/prompts/__init__.py +2 -2
  120. edsl/questions/AnswerValidatorMixin.py +289 -289
  121. edsl/questions/QuestionBase.py +664 -664
  122. edsl/questions/QuestionBaseGenMixin.py +161 -161
  123. edsl/questions/QuestionBasePromptsMixin.py +217 -217
  124. edsl/questions/QuestionBudget.py +227 -227
  125. edsl/questions/QuestionCheckBox.py +359 -359
  126. edsl/questions/QuestionExtract.py +182 -182
  127. edsl/questions/QuestionFreeText.py +114 -114
  128. edsl/questions/QuestionFunctional.py +166 -166
  129. edsl/questions/QuestionList.py +231 -231
  130. edsl/questions/QuestionMultipleChoice.py +286 -286
  131. edsl/questions/QuestionNumerical.py +153 -153
  132. edsl/questions/QuestionRank.py +324 -324
  133. edsl/questions/Quick.py +41 -41
  134. edsl/questions/RegisterQuestionsMeta.py +71 -71
  135. edsl/questions/ResponseValidatorABC.py +174 -174
  136. edsl/questions/SimpleAskMixin.py +73 -73
  137. edsl/questions/__init__.py +26 -26
  138. edsl/questions/compose_questions.py +98 -98
  139. edsl/questions/decorators.py +21 -21
  140. edsl/questions/derived/QuestionLikertFive.py +76 -76
  141. edsl/questions/derived/QuestionLinearScale.py +87 -87
  142. edsl/questions/derived/QuestionTopK.py +93 -93
  143. edsl/questions/derived/QuestionYesNo.py +82 -82
  144. edsl/questions/descriptors.py +413 -413
  145. edsl/questions/prompt_templates/question_budget.jinja +13 -13
  146. edsl/questions/prompt_templates/question_checkbox.jinja +32 -32
  147. edsl/questions/prompt_templates/question_extract.jinja +11 -11
  148. edsl/questions/prompt_templates/question_free_text.jinja +3 -3
  149. edsl/questions/prompt_templates/question_linear_scale.jinja +11 -11
  150. edsl/questions/prompt_templates/question_list.jinja +17 -17
  151. edsl/questions/prompt_templates/question_multiple_choice.jinja +33 -33
  152. edsl/questions/prompt_templates/question_numerical.jinja +36 -36
  153. edsl/questions/question_registry.py +177 -177
  154. edsl/questions/settings.py +12 -12
  155. edsl/questions/templates/budget/answering_instructions.jinja +7 -7
  156. edsl/questions/templates/budget/question_presentation.jinja +7 -7
  157. edsl/questions/templates/checkbox/answering_instructions.jinja +10 -10
  158. edsl/questions/templates/checkbox/question_presentation.jinja +22 -22
  159. edsl/questions/templates/extract/answering_instructions.jinja +7 -7
  160. edsl/questions/templates/likert_five/answering_instructions.jinja +10 -10
  161. edsl/questions/templates/likert_five/question_presentation.jinja +11 -11
  162. edsl/questions/templates/linear_scale/answering_instructions.jinja +5 -5
  163. edsl/questions/templates/linear_scale/question_presentation.jinja +5 -5
  164. edsl/questions/templates/list/answering_instructions.jinja +3 -3
  165. edsl/questions/templates/list/question_presentation.jinja +5 -5
  166. edsl/questions/templates/multiple_choice/answering_instructions.jinja +9 -9
  167. edsl/questions/templates/multiple_choice/question_presentation.jinja +11 -11
  168. edsl/questions/templates/numerical/answering_instructions.jinja +6 -6
  169. edsl/questions/templates/numerical/question_presentation.jinja +6 -6
  170. edsl/questions/templates/rank/answering_instructions.jinja +11 -11
  171. edsl/questions/templates/rank/question_presentation.jinja +15 -15
  172. edsl/questions/templates/top_k/answering_instructions.jinja +8 -8
  173. edsl/questions/templates/top_k/question_presentation.jinja +22 -22
  174. edsl/questions/templates/yes_no/answering_instructions.jinja +6 -6
  175. edsl/questions/templates/yes_no/question_presentation.jinja +11 -11
  176. edsl/results/CSSParameterizer.py +108 -108
  177. edsl/results/Dataset.py +424 -424
  178. edsl/results/DatasetExportMixin.py +731 -731
  179. edsl/results/DatasetTree.py +275 -275
  180. edsl/results/Result.py +465 -465
  181. edsl/results/Results.py +1165 -1165
  182. edsl/results/ResultsDBMixin.py +238 -238
  183. edsl/results/ResultsExportMixin.py +43 -43
  184. edsl/results/ResultsFetchMixin.py +33 -33
  185. edsl/results/ResultsGGMixin.py +121 -121
  186. edsl/results/ResultsToolsMixin.py +98 -98
  187. edsl/results/Selector.py +135 -135
  188. edsl/results/TableDisplay.py +198 -198
  189. edsl/results/__init__.py +2 -2
  190. edsl/results/table_display.css +77 -77
  191. edsl/results/tree_explore.py +115 -115
  192. edsl/scenarios/FileStore.py +632 -632
  193. edsl/scenarios/Scenario.py +601 -601
  194. edsl/scenarios/ScenarioHtmlMixin.py +64 -64
  195. edsl/scenarios/ScenarioJoin.py +127 -127
  196. edsl/scenarios/ScenarioList.py +1287 -1287
  197. edsl/scenarios/ScenarioListExportMixin.py +52 -52
  198. edsl/scenarios/ScenarioListPdfMixin.py +261 -261
  199. edsl/scenarios/__init__.py +4 -4
  200. edsl/shared.py +1 -1
  201. edsl/study/ObjectEntry.py +173 -173
  202. edsl/study/ProofOfWork.py +113 -113
  203. edsl/study/SnapShot.py +80 -80
  204. edsl/study/Study.py +528 -528
  205. edsl/study/__init__.py +4 -4
  206. edsl/surveys/DAG.py +148 -148
  207. edsl/surveys/Memory.py +31 -31
  208. edsl/surveys/MemoryPlan.py +244 -244
  209. edsl/surveys/Rule.py +326 -326
  210. edsl/surveys/RuleCollection.py +387 -387
  211. edsl/surveys/Survey.py +1801 -1801
  212. edsl/surveys/SurveyCSS.py +261 -261
  213. edsl/surveys/SurveyExportMixin.py +259 -259
  214. edsl/surveys/SurveyFlowVisualizationMixin.py +179 -179
  215. edsl/surveys/SurveyQualtricsImport.py +284 -284
  216. edsl/surveys/__init__.py +3 -3
  217. edsl/surveys/base.py +53 -53
  218. edsl/surveys/descriptors.py +56 -56
  219. edsl/surveys/instructions/ChangeInstruction.py +49 -49
  220. edsl/surveys/instructions/Instruction.py +65 -65
  221. edsl/surveys/instructions/InstructionCollection.py +77 -77
  222. edsl/templates/error_reporting/base.html +23 -23
  223. edsl/templates/error_reporting/exceptions_by_model.html +34 -34
  224. edsl/templates/error_reporting/exceptions_by_question_name.html +16 -16
  225. edsl/templates/error_reporting/exceptions_by_type.html +16 -16
  226. edsl/templates/error_reporting/interview_details.html +115 -115
  227. edsl/templates/error_reporting/interviews.html +19 -19
  228. edsl/templates/error_reporting/overview.html +4 -4
  229. edsl/templates/error_reporting/performance_plot.html +1 -1
  230. edsl/templates/error_reporting/report.css +73 -73
  231. edsl/templates/error_reporting/report.html +117 -117
  232. edsl/templates/error_reporting/report.js +25 -25
  233. edsl/tools/__init__.py +1 -1
  234. edsl/tools/clusters.py +192 -192
  235. edsl/tools/embeddings.py +27 -27
  236. edsl/tools/embeddings_plotting.py +118 -118
  237. edsl/tools/plotting.py +112 -112
  238. edsl/tools/summarize.py +18 -18
  239. edsl/utilities/SystemInfo.py +28 -28
  240. edsl/utilities/__init__.py +22 -22
  241. edsl/utilities/ast_utilities.py +25 -25
  242. edsl/utilities/data/Registry.py +6 -6
  243. edsl/utilities/data/__init__.py +1 -1
  244. edsl/utilities/data/scooter_results.json +1 -1
  245. edsl/utilities/decorators.py +77 -77
  246. edsl/utilities/gcp_bucket/cloud_storage.py +96 -96
  247. edsl/utilities/interface.py +627 -627
  248. edsl/utilities/naming_utilities.py +263 -263
  249. edsl/utilities/repair_functions.py +28 -28
  250. edsl/utilities/restricted_python.py +70 -70
  251. edsl/utilities/utilities.py +424 -424
  252. {edsl-0.1.39.dev1.dist-info → edsl-0.1.39.dev3.dist-info}/LICENSE +21 -21
  253. {edsl-0.1.39.dev1.dist-info → edsl-0.1.39.dev3.dist-info}/METADATA +1 -1
  254. edsl-0.1.39.dev3.dist-info/RECORD +277 -0
  255. edsl-0.1.39.dev1.dist-info/RECORD +0 -277
  256. {edsl-0.1.39.dev1.dist-info → edsl-0.1.39.dev3.dist-info}/WHEEL +0 -0
@@ -1,275 +1,275 @@
1
- from typing import Dict, List, Any, Optional, List
2
- from docx import Document
3
- from docx.shared import Inches, Pt
4
- from docx.enum.text import WD_ALIGN_PARAGRAPH
5
- from docx.enum.style import WD_STYLE_TYPE
6
-
7
-
8
- def is_hashable(v):
9
- try:
10
- hash(v)
11
- return True
12
- except TypeError:
13
- return False
14
-
15
-
16
- class TreeNode:
17
- def __init__(self, key=None, value=None):
18
- self.key = key
19
- self.value = value
20
- self.children = {}
21
-
22
-
23
- class Tree:
24
- def __init__(self, data: "Dataset", node_order: Optional[List[str]] = None):
25
- d = {}
26
- for entry in data:
27
- d.update(entry)
28
- self.data = d
29
- self.root = None
30
-
31
- self.node_order = node_order
32
-
33
- self.construct_tree(node_order)
34
-
35
- def unique_values_by_keys(self) -> dict:
36
- unique_values = {}
37
- for key, raw_values in self.data.items():
38
- values = [v if is_hashable(v) else str(v) for v in raw_values]
39
- unique_values[key] = list(set(values))
40
- return unique_values
41
-
42
- def construct_tree(self, node_order: Optional[List[str]] = None):
43
- # Validate node_order
44
- if node_order is None:
45
- unique_values = self.unique_values_by_keys()
46
- # Sort keys by number of unique values
47
- node_order = sorted(
48
- unique_values, key=lambda k: len(unique_values[k]), reverse=True
49
- )
50
- else:
51
- if not set(node_order).issubset(set(self.data.keys())):
52
- invalid_keys = set(node_order) - set(self.data.keys())
53
- raise ValueError(f"Invalid keys in node_order: {invalid_keys}")
54
-
55
- self.root = TreeNode()
56
-
57
- for i in range(len(self.data[node_order[0]])):
58
- current = self.root
59
- for level in node_order[:-1]:
60
- value = self.data[level][i]
61
- if not is_hashable(value):
62
- value = str(value)
63
- if value not in current.children:
64
- current.children[value] = TreeNode(key=level, value=value)
65
- current = current.children[value]
66
-
67
- leaf_key = node_order[-1]
68
- leaf_value = self.data[leaf_key][i]
69
- if not is_hashable(leaf_value):
70
- leaf_value = str(leaf_value)
71
- if leaf_value not in current.children:
72
- current.children[leaf_value] = TreeNode(key=leaf_key, value=leaf_value)
73
-
74
- def __repr__(self):
75
- if self.node_order is not None:
76
- return f"Tree(Dataset({self.data}), node_order={self.node_order})"
77
- else:
78
- return f"Tree(Dataset({self.data}))"
79
-
80
- def print_tree(
81
- self, node: Optional[TreeNode] = None, level: int = 0, print_keys: bool = False
82
- ):
83
- if node is None:
84
- node = self.root
85
- if node is None:
86
- print("Tree has not been constructed yet.")
87
- return
88
-
89
- if node.value is not None:
90
- if print_keys and node.key is not None:
91
- print(" " * level + f"{node.key}: {node.value}")
92
- else:
93
- print(" " * level + str(node.value))
94
- for child in node.children.values():
95
- self.print_tree(child, level + 1, print_keys)
96
-
97
- def to_docx(self, filename: Optional[str] = None):
98
- if filename is None:
99
- filename = "tree_structure.docx"
100
-
101
- doc = Document()
102
-
103
- # Create styles for headings
104
- for i in range(1, 10): # Up to 9 levels of headings
105
- style_name = f"Heading {i}"
106
- if style_name not in doc.styles:
107
- doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
108
-
109
- # Get or create the 'Body Text' style
110
- if "Body Text" not in doc.styles:
111
- body_style = doc.styles.add_style("Body Text", WD_STYLE_TYPE.PARAGRAPH)
112
- else:
113
- body_style = doc.styles["Body Text"]
114
-
115
- body_style.font.size = Pt(11)
116
-
117
- self._add_to_docx(doc, self.root, 0)
118
- doc.save(filename)
119
- from edsl.utilities.utilities import file_notice
120
-
121
- file_notice(filename)
122
-
123
- def _repr_html_(self):
124
- """Returns an interactive HTML representation of the tree with collapsible sections."""
125
-
126
- # Generate a unique ID for this tree instance
127
- import uuid
128
-
129
- tree_id = f"tree_{uuid.uuid4().hex[:8]}"
130
-
131
- styles = f"""
132
- <div class="{tree_id}">
133
- <style>
134
- .{tree_id} details {{
135
- margin-left: 20px;
136
- }}
137
- .{tree_id} summary {{
138
- cursor: pointer;
139
- margin: 2px 0;
140
- }}
141
- .{tree_id} .value {{
142
- font-family: monospace;
143
- background: #f5f5f5;
144
- padding: 2px 6px;
145
- border-radius: 3px;
146
- margin: 1px 0;
147
- }}
148
- .{tree_id} .key {{
149
- color: #666;
150
- font-style: italic;
151
- }}
152
- </style>
153
- """
154
-
155
- def node_to_html(node, level=0, print_keys=True):
156
- if node is None:
157
- return "Tree has not been constructed yet."
158
-
159
- html = []
160
-
161
- if node.value is not None:
162
- # Create the node content
163
- content = []
164
- if print_keys and node.key is not None:
165
- content.append(f'<span class="key">{node.key}: </span>')
166
- content.append(f'<span class="value">{node.value}</span>')
167
- content_html = "".join(content)
168
-
169
- if node.children:
170
- # Node with children
171
- html.append(f'<details {"open" if level < 1 else ""}>')
172
- html.append(f"<summary>{content_html}</summary>")
173
- for child in node.children.values():
174
- html.append(node_to_html(child, level + 1, print_keys))
175
- html.append("</details>")
176
- else:
177
- # Leaf node
178
- html.append(f"<div>{content_html}</div>")
179
- else:
180
- # Root node with no value
181
- if node.children:
182
- for child in node.children.values():
183
- html.append(node_to_html(child, level, print_keys))
184
-
185
- return "\n".join(html)
186
-
187
- tree_html = node_to_html(self.root)
188
- return f"{styles}{tree_html}</div>"
189
-
190
- # def _repr_html_(self):
191
- # """Returns an HTML representation of the tree, following the same logic as print_tree."""
192
- # styles = """
193
- # <style>
194
- # .tree-container {
195
- # font-family: monospace;
196
- # white-space: pre;
197
- # margin: 10px;
198
- # }
199
- # </style>
200
- # """
201
-
202
- # def node_to_html(node, level=0, print_keys=False):
203
- # if node is None:
204
- # node = self.root
205
- # if node is None:
206
- # return "Tree has not been constructed yet."
207
-
208
- # html = []
209
- # if node.value is not None:
210
- # indent = "&nbsp;" * 2 * level # Using &nbsp; for HTML spaces
211
- # if print_keys and node.key is not None:
212
- # html.append(f"{indent}{node.key}: {node.value}<br>")
213
- # else:
214
- # html.append(f"{indent}{node.value}<br>")
215
-
216
- # for child in node.children.values():
217
- # html.append(node_to_html(child, level + 1, print_keys))
218
-
219
- # return "".join(html)
220
-
221
- # tree_html = node_to_html(self.root)
222
- # return f'<div class="tree-container">{tree_html}</div>{styles}'
223
-
224
- def _add_to_docx(self, doc, node: TreeNode, level: int):
225
- if node.value is not None:
226
- if level == 0:
227
- doc.add_heading(str(node.value), level=level + 1)
228
- elif node.children: # If the node has children, it's not the last level
229
- para = doc.add_paragraph(str(node.value))
230
- para.style = f"Heading {level+1}"
231
- else: # If the node has no children, it's the last level (body text)
232
- para = doc.add_paragraph(str(node.value))
233
- para.style = "Body Text"
234
-
235
- # Process child nodes (moved outside the if block)
236
- for child in node.children.values():
237
- self._add_to_docx(doc, child, level + 1)
238
-
239
-
240
- # Example usage (commented out)
241
- """
242
- from edsl.results.Dataset import Dataset
243
-
244
- data = Dataset(
245
- [
246
- {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
247
- {"country": ["US", "China", "France", "Canada", "Japan"]},
248
- {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
249
- {"population": [8419000, 21540000, 2161000, 2930000, 13960000]},
250
- ]
251
- )
252
-
253
- tree = Tree(data)
254
-
255
- try:
256
- tree.construct_tree(["continent", "country", "city", "population"])
257
- print("Tree without key names:")
258
- tree.print_tree()
259
- print("\nTree with key names:")
260
- tree.print_tree(print_keys=True)
261
- except ValueError as e:
262
- print(f"Error: {e}")
263
-
264
- # Demonstrating validation
265
- try:
266
- tree.construct_tree(["continent", "country", "invalid_key"])
267
- except ValueError as e:
268
- print(f"\nValidation Error: {e}")
269
-
270
- tree = Tree(data)
271
- tree.construct_tree(["continent", "country", "city", "population"])
272
- tree.print_tree(print_keys=True)
273
- tree.to_docx("tree_structure.docx")
274
- print("DocX file 'tree_structure.docx' has been created.")
275
- """
1
+ from typing import Dict, List, Any, Optional, List
2
+ from docx import Document
3
+ from docx.shared import Inches, Pt
4
+ from docx.enum.text import WD_ALIGN_PARAGRAPH
5
+ from docx.enum.style import WD_STYLE_TYPE
6
+
7
+
8
+ def is_hashable(v):
9
+ try:
10
+ hash(v)
11
+ return True
12
+ except TypeError:
13
+ return False
14
+
15
+
16
+ class TreeNode:
17
+ def __init__(self, key=None, value=None):
18
+ self.key = key
19
+ self.value = value
20
+ self.children = {}
21
+
22
+
23
+ class Tree:
24
+ def __init__(self, data: "Dataset", node_order: Optional[List[str]] = None):
25
+ d = {}
26
+ for entry in data:
27
+ d.update(entry)
28
+ self.data = d
29
+ self.root = None
30
+
31
+ self.node_order = node_order
32
+
33
+ self.construct_tree(node_order)
34
+
35
+ def unique_values_by_keys(self) -> dict:
36
+ unique_values = {}
37
+ for key, raw_values in self.data.items():
38
+ values = [v if is_hashable(v) else str(v) for v in raw_values]
39
+ unique_values[key] = list(set(values))
40
+ return unique_values
41
+
42
+ def construct_tree(self, node_order: Optional[List[str]] = None):
43
+ # Validate node_order
44
+ if node_order is None:
45
+ unique_values = self.unique_values_by_keys()
46
+ # Sort keys by number of unique values
47
+ node_order = sorted(
48
+ unique_values, key=lambda k: len(unique_values[k]), reverse=True
49
+ )
50
+ else:
51
+ if not set(node_order).issubset(set(self.data.keys())):
52
+ invalid_keys = set(node_order) - set(self.data.keys())
53
+ raise ValueError(f"Invalid keys in node_order: {invalid_keys}")
54
+
55
+ self.root = TreeNode()
56
+
57
+ for i in range(len(self.data[node_order[0]])):
58
+ current = self.root
59
+ for level in node_order[:-1]:
60
+ value = self.data[level][i]
61
+ if not is_hashable(value):
62
+ value = str(value)
63
+ if value not in current.children:
64
+ current.children[value] = TreeNode(key=level, value=value)
65
+ current = current.children[value]
66
+
67
+ leaf_key = node_order[-1]
68
+ leaf_value = self.data[leaf_key][i]
69
+ if not is_hashable(leaf_value):
70
+ leaf_value = str(leaf_value)
71
+ if leaf_value not in current.children:
72
+ current.children[leaf_value] = TreeNode(key=leaf_key, value=leaf_value)
73
+
74
+ def __repr__(self):
75
+ if self.node_order is not None:
76
+ return f"Tree(Dataset({self.data}), node_order={self.node_order})"
77
+ else:
78
+ return f"Tree(Dataset({self.data}))"
79
+
80
+ def print_tree(
81
+ self, node: Optional[TreeNode] = None, level: int = 0, print_keys: bool = False
82
+ ):
83
+ if node is None:
84
+ node = self.root
85
+ if node is None:
86
+ print("Tree has not been constructed yet.")
87
+ return
88
+
89
+ if node.value is not None:
90
+ if print_keys and node.key is not None:
91
+ print(" " * level + f"{node.key}: {node.value}")
92
+ else:
93
+ print(" " * level + str(node.value))
94
+ for child in node.children.values():
95
+ self.print_tree(child, level + 1, print_keys)
96
+
97
+ def to_docx(self, filename: Optional[str] = None):
98
+ if filename is None:
99
+ filename = "tree_structure.docx"
100
+
101
+ doc = Document()
102
+
103
+ # Create styles for headings
104
+ for i in range(1, 10): # Up to 9 levels of headings
105
+ style_name = f"Heading {i}"
106
+ if style_name not in doc.styles:
107
+ doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
108
+
109
+ # Get or create the 'Body Text' style
110
+ if "Body Text" not in doc.styles:
111
+ body_style = doc.styles.add_style("Body Text", WD_STYLE_TYPE.PARAGRAPH)
112
+ else:
113
+ body_style = doc.styles["Body Text"]
114
+
115
+ body_style.font.size = Pt(11)
116
+
117
+ self._add_to_docx(doc, self.root, 0)
118
+ doc.save(filename)
119
+ from edsl.utilities.utilities import file_notice
120
+
121
+ file_notice(filename)
122
+
123
+ def _repr_html_(self):
124
+ """Returns an interactive HTML representation of the tree with collapsible sections."""
125
+
126
+ # Generate a unique ID for this tree instance
127
+ import uuid
128
+
129
+ tree_id = f"tree_{uuid.uuid4().hex[:8]}"
130
+
131
+ styles = f"""
132
+ <div class="{tree_id}">
133
+ <style>
134
+ .{tree_id} details {{
135
+ margin-left: 20px;
136
+ }}
137
+ .{tree_id} summary {{
138
+ cursor: pointer;
139
+ margin: 2px 0;
140
+ }}
141
+ .{tree_id} .value {{
142
+ font-family: monospace;
143
+ background: #f5f5f5;
144
+ padding: 2px 6px;
145
+ border-radius: 3px;
146
+ margin: 1px 0;
147
+ }}
148
+ .{tree_id} .key {{
149
+ color: #666;
150
+ font-style: italic;
151
+ }}
152
+ </style>
153
+ """
154
+
155
+ def node_to_html(node, level=0, print_keys=True):
156
+ if node is None:
157
+ return "Tree has not been constructed yet."
158
+
159
+ html = []
160
+
161
+ if node.value is not None:
162
+ # Create the node content
163
+ content = []
164
+ if print_keys and node.key is not None:
165
+ content.append(f'<span class="key">{node.key}: </span>')
166
+ content.append(f'<span class="value">{node.value}</span>')
167
+ content_html = "".join(content)
168
+
169
+ if node.children:
170
+ # Node with children
171
+ html.append(f'<details {"open" if level < 1 else ""}>')
172
+ html.append(f"<summary>{content_html}</summary>")
173
+ for child in node.children.values():
174
+ html.append(node_to_html(child, level + 1, print_keys))
175
+ html.append("</details>")
176
+ else:
177
+ # Leaf node
178
+ html.append(f"<div>{content_html}</div>")
179
+ else:
180
+ # Root node with no value
181
+ if node.children:
182
+ for child in node.children.values():
183
+ html.append(node_to_html(child, level, print_keys))
184
+
185
+ return "\n".join(html)
186
+
187
+ tree_html = node_to_html(self.root)
188
+ return f"{styles}{tree_html}</div>"
189
+
190
+ # def _repr_html_(self):
191
+ # """Returns an HTML representation of the tree, following the same logic as print_tree."""
192
+ # styles = """
193
+ # <style>
194
+ # .tree-container {
195
+ # font-family: monospace;
196
+ # white-space: pre;
197
+ # margin: 10px;
198
+ # }
199
+ # </style>
200
+ # """
201
+
202
+ # def node_to_html(node, level=0, print_keys=False):
203
+ # if node is None:
204
+ # node = self.root
205
+ # if node is None:
206
+ # return "Tree has not been constructed yet."
207
+
208
+ # html = []
209
+ # if node.value is not None:
210
+ # indent = "&nbsp;" * 2 * level # Using &nbsp; for HTML spaces
211
+ # if print_keys and node.key is not None:
212
+ # html.append(f"{indent}{node.key}: {node.value}<br>")
213
+ # else:
214
+ # html.append(f"{indent}{node.value}<br>")
215
+
216
+ # for child in node.children.values():
217
+ # html.append(node_to_html(child, level + 1, print_keys))
218
+
219
+ # return "".join(html)
220
+
221
+ # tree_html = node_to_html(self.root)
222
+ # return f'<div class="tree-container">{tree_html}</div>{styles}'
223
+
224
+ def _add_to_docx(self, doc, node: TreeNode, level: int):
225
+ if node.value is not None:
226
+ if level == 0:
227
+ doc.add_heading(str(node.value), level=level + 1)
228
+ elif node.children: # If the node has children, it's not the last level
229
+ para = doc.add_paragraph(str(node.value))
230
+ para.style = f"Heading {level+1}"
231
+ else: # If the node has no children, it's the last level (body text)
232
+ para = doc.add_paragraph(str(node.value))
233
+ para.style = "Body Text"
234
+
235
+ # Process child nodes (moved outside the if block)
236
+ for child in node.children.values():
237
+ self._add_to_docx(doc, child, level + 1)
238
+
239
+
240
+ # Example usage (commented out)
241
+ """
242
+ from edsl.results.Dataset import Dataset
243
+
244
+ data = Dataset(
245
+ [
246
+ {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
247
+ {"country": ["US", "China", "France", "Canada", "Japan"]},
248
+ {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
249
+ {"population": [8419000, 21540000, 2161000, 2930000, 13960000]},
250
+ ]
251
+ )
252
+
253
+ tree = Tree(data)
254
+
255
+ try:
256
+ tree.construct_tree(["continent", "country", "city", "population"])
257
+ print("Tree without key names:")
258
+ tree.print_tree()
259
+ print("\nTree with key names:")
260
+ tree.print_tree(print_keys=True)
261
+ except ValueError as e:
262
+ print(f"Error: {e}")
263
+
264
+ # Demonstrating validation
265
+ try:
266
+ tree.construct_tree(["continent", "country", "invalid_key"])
267
+ except ValueError as e:
268
+ print(f"\nValidation Error: {e}")
269
+
270
+ tree = Tree(data)
271
+ tree.construct_tree(["continent", "country", "city", "population"])
272
+ tree.print_tree(print_keys=True)
273
+ tree.to_docx("tree_structure.docx")
274
+ print("DocX file 'tree_structure.docx' has been created.")
275
+ """