edsl 0.1.37.dev6__py3-none-any.whl → 0.1.38__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 (271) hide show
  1. edsl/Base.py +332 -303
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +49 -48
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +867 -855
  7. edsl/agents/AgentList.py +413 -350
  8. edsl/agents/Invigilator.py +233 -222
  9. edsl/agents/InvigilatorBase.py +265 -284
  10. edsl/agents/PromptConstructor.py +354 -353
  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 -289
  26. edsl/config.py +157 -149
  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 -958
  35. edsl/coop/utils.py +131 -131
  36. edsl/data/Cache.py +555 -527
  37. edsl/data/CacheEntry.py +233 -228
  38. edsl/data/CacheHandler.py +149 -149
  39. edsl/data/RemoteCacheSync.py +78 -97
  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 -173
  45. edsl/exceptions/BaseException.py +21 -21
  46. edsl/exceptions/__init__.py +54 -54
  47. edsl/exceptions/agents.py +42 -38
  48. edsl/exceptions/cache.py +5 -0
  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 -156
  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 -0
  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 -39
  77. edsl/inference_services/write_available.py +10 -10
  78. edsl/jobs/Answers.py +56 -56
  79. edsl/jobs/Jobs.py +898 -1347
  80. edsl/jobs/JobsChecks.py +147 -0
  81. edsl/jobs/JobsPrompts.py +268 -0
  82. edsl/jobs/JobsRemoteInferenceHandler.py +239 -0
  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 -248
  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 -338
  97. edsl/jobs/runners/JobsRunnerStatus.py +330 -332
  98. edsl/jobs/tasks/QuestionTaskCreator.py +242 -242
  99. edsl/jobs/tasks/TaskCreators.py +64 -64
  100. edsl/jobs/tasks/TaskHistory.py +450 -442
  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 -706
  107. edsl/language_models/ModelList.py +155 -102
  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 -137
  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 -259
  117. edsl/notebooks/__init__.py +1 -1
  118. edsl/prompts/Prompt.py +362 -357
  119. edsl/prompts/__init__.py +2 -2
  120. edsl/questions/AnswerValidatorMixin.py +289 -289
  121. edsl/questions/QuestionBase.py +664 -656
  122. edsl/questions/QuestionBaseGenMixin.py +161 -161
  123. edsl/questions/QuestionBasePromptsMixin.py +217 -234
  124. edsl/questions/QuestionBudget.py +227 -227
  125. edsl/questions/QuestionCheckBox.py +359 -359
  126. edsl/questions/QuestionExtract.py +182 -183
  127. edsl/questions/QuestionFreeText.py +114 -114
  128. edsl/questions/QuestionFunctional.py +166 -159
  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 -91
  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 -147
  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 -0
  177. edsl/results/Dataset.py +424 -293
  178. edsl/results/DatasetExportMixin.py +731 -717
  179. edsl/results/DatasetTree.py +275 -145
  180. edsl/results/Result.py +465 -450
  181. edsl/results/Results.py +1165 -1071
  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 -0
  189. edsl/results/__init__.py +2 -2
  190. edsl/results/table_display.css +78 -0
  191. edsl/results/tree_explore.py +115 -115
  192. edsl/scenarios/FileStore.py +632 -458
  193. edsl/scenarios/Scenario.py +601 -546
  194. edsl/scenarios/ScenarioHtmlMixin.py +64 -64
  195. edsl/scenarios/ScenarioJoin.py +127 -0
  196. edsl/scenarios/ScenarioList.py +1287 -1112
  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 -330
  210. edsl/surveys/RuleCollection.py +387 -387
  211. edsl/surveys/Survey.py +1801 -1795
  212. edsl/surveys/SurveyCSS.py +261 -261
  213. edsl/surveys/SurveyExportMixin.py +259 -259
  214. edsl/surveys/SurveyFlowVisualizationMixin.py +179 -121
  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 -47
  220. edsl/surveys/instructions/Instruction.py +65 -51
  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 -10
  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/{conjure → 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 -409
  252. {edsl-0.1.37.dev6.dist-info → edsl-0.1.38.dist-info}/LICENSE +21 -21
  253. {edsl-0.1.37.dev6.dist-info → edsl-0.1.38.dist-info}/METADATA +2 -1
  254. edsl-0.1.38.dist-info/RECORD +277 -0
  255. edsl/conjure/AgentConstructionMixin.py +0 -160
  256. edsl/conjure/Conjure.py +0 -62
  257. edsl/conjure/InputData.py +0 -659
  258. edsl/conjure/InputDataCSV.py +0 -48
  259. edsl/conjure/InputDataMixinQuestionStats.py +0 -182
  260. edsl/conjure/InputDataPyRead.py +0 -91
  261. edsl/conjure/InputDataSPSS.py +0 -8
  262. edsl/conjure/InputDataStata.py +0 -8
  263. edsl/conjure/QuestionOptionMixin.py +0 -76
  264. edsl/conjure/QuestionTypeMixin.py +0 -23
  265. edsl/conjure/RawQuestion.py +0 -65
  266. edsl/conjure/SurveyResponses.py +0 -7
  267. edsl/conjure/__init__.py +0 -9
  268. edsl/conjure/examples/placeholder.txt +0 -0
  269. edsl/conjure/utilities.py +0 -201
  270. edsl-0.1.37.dev6.dist-info/RECORD +0 -283
  271. {edsl-0.1.37.dev6.dist-info → edsl-0.1.38.dist-info}/WHEEL +0 -0
@@ -1,238 +1,238 @@
1
- """Mixin for working with SQLite respresentation of a 'Results' object."""
2
-
3
- import sqlite3
4
- from enum import Enum
5
- from typing import Literal, Union, Optional
6
-
7
-
8
- class SQLDataShape(Enum):
9
- """Enum for the shape of the data in the SQL database."""
10
-
11
- WIDE = "wide"
12
- LONG = "long"
13
-
14
-
15
- class ResultsDBMixin:
16
- """Mixin for interacting with a Results object as if it were a SQL database."""
17
-
18
- def _rows(self):
19
- """Return the rows of the `Results` object as a list of tuples."""
20
- for index, result in enumerate(self):
21
- yield from result.rows(index)
22
-
23
- def export_sql_dump(self, shape: Literal["wide", "long"], filename: str):
24
- """Export the SQL database to a file.
25
-
26
- :param shape: The shape of the data in the database (wide or long)
27
- :param filename: The filename to save the database to
28
- """
29
- shape_enum = self._get_shape_enum(shape)
30
- conn = self._db(shape=shape_enum)
31
-
32
- with open(filename, "w") as f:
33
- for line in conn.iterdump():
34
- f.write(f"{line}\n")
35
-
36
- conn.close()
37
-
38
- def backup_db_to_file(self, shape: Literal["wide", "long"], filename: str):
39
- """Backup the in-memory database to a file.
40
-
41
-
42
- :param shape: The shape of the data in the database (wide or long)
43
- :param filename: The filename to save the database to
44
-
45
- >>> from edsl.results import Results
46
- >>> r = Results.example()
47
- >>> r.backup_db_to_file(filename="backup.db", shape="long")
48
-
49
- """
50
- shape_enum = self._get_shape_enum(shape)
51
- # Source database connection (in-memory)
52
- source_conn = self._db(shape=shape_enum)
53
-
54
- # Destination database connection (file)
55
- dest_conn = sqlite3.connect(filename)
56
-
57
- # Backup in-memory database to file
58
- with source_conn:
59
- source_conn.backup(dest_conn)
60
-
61
- # Close both connections
62
- source_conn.close()
63
- dest_conn.close()
64
-
65
- def _db(self, shape: SQLDataShape, remove_prefix=False):
66
- """Create a SQLite database in memory and return the connection.
67
-
68
- :param shape: The shape of the data in the database (wide or long)
69
- :param remove_prefix: Whether to remove the prefix from the column names
70
-
71
- """
72
- if shape == SQLDataShape.LONG:
73
- conn = sqlite3.connect(":memory:")
74
-
75
- create_table_query = """
76
- CREATE TABLE self (
77
- id INTEGER,
78
- data_type TEXT,
79
- key TEXT,
80
- value TEXT
81
- )
82
- """
83
- conn.execute(create_table_query)
84
-
85
- list_of_tuples = list(self._rows())
86
- insert_query = (
87
- "INSERT INTO self (id, data_type, key, value) VALUES (?, ?, ?, ?)"
88
- )
89
- conn.executemany(insert_query, list_of_tuples)
90
- conn.commit()
91
- return conn
92
- elif shape == SQLDataShape.WIDE:
93
- from sqlalchemy import create_engine
94
-
95
- engine = create_engine("sqlite:///:memory:")
96
- df = self.to_pandas(remove_prefix=remove_prefix, lists_as_strings=True)
97
- df.to_sql("self", engine, index=False, if_exists="replace")
98
- return engine.connect()
99
- else:
100
- raise Exception("Invalid SQLDataShape")
101
-
102
- def _get_shape_enum(self, shape: Literal["wide", "long"]):
103
- """Convert the shape string to a SQLDataShape enum."""
104
- if shape is None:
105
- raise Exception("Must select either 'wide' or 'long' format")
106
- elif shape == "wide":
107
- return SQLDataShape.WIDE
108
- elif shape == "long":
109
- return SQLDataShape.LONG
110
- else:
111
- raise Exception("Invalid shape: must be either 'long' or 'wide'")
112
-
113
- def sql(
114
- self,
115
- query: str,
116
- shape: Literal["wide", "long"] = "wide",
117
- remove_prefix: bool = True,
118
- transpose: bool = None,
119
- transpose_by: str = None,
120
- csv: bool = False,
121
- to_list=False,
122
- to_latex=False,
123
- filename: Optional[str] = None,
124
- ) -> Union["pd.DataFrame", str]:
125
- """Execute a SQL query and return the results as a DataFrame.
126
-
127
- :param query: The SQL query to execute
128
- :param shape: The shape of the data in the database (wide or long)
129
- :param remove_prefix: Whether to remove the prefix from the column names
130
- :param transpose: Whether to transpose the DataFrame
131
- :param transpose_by: The column to use as the index when transposing
132
- :param csv: Whether to return the DataFrame as a CSV string
133
-
134
-
135
- Example usage:
136
-
137
- >>> from edsl.results import Results
138
- >>> r = Results.example()
139
- >>> d = r.sql("select data_type, key, value from self where data_type = 'answer' order by value limit 3", shape="long")
140
- >>> sorted(list(d['value']))
141
- ['Good', 'Great', 'Great']
142
-
143
- We can also return the data in wide format.
144
- Note the use of single quotes to escape the column names, as required by sql.
145
-
146
- >>> from edsl.results import Results
147
- >>> Results.example().sql("select how_feeling from self", shape = 'wide', remove_prefix=True)
148
- how_feeling
149
- 0 OK
150
- 1 Great
151
- 2 Terrible
152
- 3 OK
153
- """
154
- import pandas as pd
155
-
156
- shape_enum = self._get_shape_enum(shape)
157
-
158
- conn = self._db(shape=shape_enum, remove_prefix=remove_prefix)
159
- df = pd.read_sql_query(query, conn)
160
-
161
- # Transpose the DataFrame if transpose is True
162
- if transpose or transpose_by:
163
- df = pd.DataFrame(df)
164
- if transpose_by:
165
- df = df.set_index(transpose_by)
166
- else:
167
- df = df.set_index(df.columns[0])
168
- df = df.transpose()
169
-
170
- if csv and to_list:
171
- raise Exception("Cannot return both CSV and list")
172
-
173
- if to_list:
174
- return df.values.tolist()
175
-
176
- if to_latex:
177
- df.columns = [col.replace("_", " ") for col in df.columns]
178
-
179
- latex_output = df.to_latex(index=False)
180
- if filename:
181
- with open(filename, "w") as f:
182
- f.write(latex_output)
183
- return None
184
- return latex_output
185
-
186
- if csv:
187
- if filename:
188
- df.to_csv(filename, index=False)
189
- return None
190
-
191
- return df.to_csv(index=False)
192
-
193
- return df
194
-
195
- def show_schema(
196
- self, shape: Literal["wide", "long"], remove_prefix: bool = False
197
- ) -> None:
198
- """Show the schema of the Results database.
199
-
200
- :param shape: The shape of the data in the database (wide or long)
201
- :param remove_prefix: Whether to remove the prefix from the column names
202
-
203
- >>> from edsl.results import Results
204
- >>> r = Results.example()
205
- >>> r.show_schema(shape="long")
206
- Type: table, Name: self, SQL: CREATE TABLE self (
207
- ...
208
- <BLANKLINE>
209
- """
210
- import pandas as pd
211
-
212
- shape_enum = self._get_shape_enum(shape)
213
- conn = self._db(shape=shape_enum, remove_prefix=remove_prefix)
214
-
215
- if shape_enum == SQLDataShape.LONG:
216
- # Query to get the schema of all tables
217
- query = "SELECT type, name, sql FROM sqlite_master WHERE type='table'"
218
- cursor = conn.execute(query)
219
- schema = cursor.fetchall()
220
- conn.close()
221
-
222
- # Format and return the schema information
223
- schema_info = ""
224
- for row in schema:
225
- schema_info += f"Type: {row[0]}, Name: {row[1]}, SQL: {row[2]}\n"
226
-
227
- print(schema_info)
228
- elif shape_enum == SQLDataShape.WIDE:
229
- query = f"PRAGMA table_info(self)"
230
- schema = pd.read_sql(query, conn)
231
- # print(schema)
232
- return schema
233
-
234
-
235
- if __name__ == "__main__":
236
- import doctest
237
-
238
- doctest.testmod(optionflags=doctest.ELLIPSIS)
1
+ """Mixin for working with SQLite respresentation of a 'Results' object."""
2
+
3
+ import sqlite3
4
+ from enum import Enum
5
+ from typing import Literal, Union, Optional
6
+
7
+
8
+ class SQLDataShape(Enum):
9
+ """Enum for the shape of the data in the SQL database."""
10
+
11
+ WIDE = "wide"
12
+ LONG = "long"
13
+
14
+
15
+ class ResultsDBMixin:
16
+ """Mixin for interacting with a Results object as if it were a SQL database."""
17
+
18
+ def _rows(self):
19
+ """Return the rows of the `Results` object as a list of tuples."""
20
+ for index, result in enumerate(self):
21
+ yield from result.rows(index)
22
+
23
+ def export_sql_dump(self, shape: Literal["wide", "long"], filename: str):
24
+ """Export the SQL database to a file.
25
+
26
+ :param shape: The shape of the data in the database (wide or long)
27
+ :param filename: The filename to save the database to
28
+ """
29
+ shape_enum = self._get_shape_enum(shape)
30
+ conn = self._db(shape=shape_enum)
31
+
32
+ with open(filename, "w") as f:
33
+ for line in conn.iterdump():
34
+ f.write(f"{line}\n")
35
+
36
+ conn.close()
37
+
38
+ def backup_db_to_file(self, shape: Literal["wide", "long"], filename: str):
39
+ """Backup the in-memory database to a file.
40
+
41
+
42
+ :param shape: The shape of the data in the database (wide or long)
43
+ :param filename: The filename to save the database to
44
+
45
+ >>> from edsl.results import Results
46
+ >>> r = Results.example()
47
+ >>> r.backup_db_to_file(filename="backup.db", shape="long")
48
+
49
+ """
50
+ shape_enum = self._get_shape_enum(shape)
51
+ # Source database connection (in-memory)
52
+ source_conn = self._db(shape=shape_enum)
53
+
54
+ # Destination database connection (file)
55
+ dest_conn = sqlite3.connect(filename)
56
+
57
+ # Backup in-memory database to file
58
+ with source_conn:
59
+ source_conn.backup(dest_conn)
60
+
61
+ # Close both connections
62
+ source_conn.close()
63
+ dest_conn.close()
64
+
65
+ def _db(self, shape: SQLDataShape, remove_prefix=False):
66
+ """Create a SQLite database in memory and return the connection.
67
+
68
+ :param shape: The shape of the data in the database (wide or long)
69
+ :param remove_prefix: Whether to remove the prefix from the column names
70
+
71
+ """
72
+ if shape == SQLDataShape.LONG:
73
+ conn = sqlite3.connect(":memory:")
74
+
75
+ create_table_query = """
76
+ CREATE TABLE self (
77
+ id INTEGER,
78
+ data_type TEXT,
79
+ key TEXT,
80
+ value TEXT
81
+ )
82
+ """
83
+ conn.execute(create_table_query)
84
+
85
+ list_of_tuples = list(self._rows())
86
+ insert_query = (
87
+ "INSERT INTO self (id, data_type, key, value) VALUES (?, ?, ?, ?)"
88
+ )
89
+ conn.executemany(insert_query, list_of_tuples)
90
+ conn.commit()
91
+ return conn
92
+ elif shape == SQLDataShape.WIDE:
93
+ from sqlalchemy import create_engine
94
+
95
+ engine = create_engine("sqlite:///:memory:")
96
+ df = self.to_pandas(remove_prefix=remove_prefix, lists_as_strings=True)
97
+ df.to_sql("self", engine, index=False, if_exists="replace")
98
+ return engine.connect()
99
+ else:
100
+ raise Exception("Invalid SQLDataShape")
101
+
102
+ def _get_shape_enum(self, shape: Literal["wide", "long"]):
103
+ """Convert the shape string to a SQLDataShape enum."""
104
+ if shape is None:
105
+ raise Exception("Must select either 'wide' or 'long' format")
106
+ elif shape == "wide":
107
+ return SQLDataShape.WIDE
108
+ elif shape == "long":
109
+ return SQLDataShape.LONG
110
+ else:
111
+ raise Exception("Invalid shape: must be either 'long' or 'wide'")
112
+
113
+ def sql(
114
+ self,
115
+ query: str,
116
+ shape: Literal["wide", "long"] = "wide",
117
+ remove_prefix: bool = True,
118
+ transpose: bool = None,
119
+ transpose_by: str = None,
120
+ csv: bool = False,
121
+ to_list=False,
122
+ to_latex=False,
123
+ filename: Optional[str] = None,
124
+ ) -> Union["pd.DataFrame", str]:
125
+ """Execute a SQL query and return the results as a DataFrame.
126
+
127
+ :param query: The SQL query to execute
128
+ :param shape: The shape of the data in the database (wide or long)
129
+ :param remove_prefix: Whether to remove the prefix from the column names
130
+ :param transpose: Whether to transpose the DataFrame
131
+ :param transpose_by: The column to use as the index when transposing
132
+ :param csv: Whether to return the DataFrame as a CSV string
133
+
134
+
135
+ Example usage:
136
+
137
+ >>> from edsl.results import Results
138
+ >>> r = Results.example()
139
+ >>> d = r.sql("select data_type, key, value from self where data_type = 'answer' order by value limit 3", shape="long")
140
+ >>> sorted(list(d['value']))
141
+ ['Good', 'Great', 'Great']
142
+
143
+ We can also return the data in wide format.
144
+ Note the use of single quotes to escape the column names, as required by sql.
145
+
146
+ >>> from edsl.results import Results
147
+ >>> Results.example().sql("select how_feeling from self", shape = 'wide', remove_prefix=True)
148
+ how_feeling
149
+ 0 OK
150
+ 1 Great
151
+ 2 Terrible
152
+ 3 OK
153
+ """
154
+ import pandas as pd
155
+
156
+ shape_enum = self._get_shape_enum(shape)
157
+
158
+ conn = self._db(shape=shape_enum, remove_prefix=remove_prefix)
159
+ df = pd.read_sql_query(query, conn)
160
+
161
+ # Transpose the DataFrame if transpose is True
162
+ if transpose or transpose_by:
163
+ df = pd.DataFrame(df)
164
+ if transpose_by:
165
+ df = df.set_index(transpose_by)
166
+ else:
167
+ df = df.set_index(df.columns[0])
168
+ df = df.transpose()
169
+
170
+ if csv and to_list:
171
+ raise Exception("Cannot return both CSV and list")
172
+
173
+ if to_list:
174
+ return df.values.tolist()
175
+
176
+ if to_latex:
177
+ df.columns = [col.replace("_", " ") for col in df.columns]
178
+
179
+ latex_output = df.to_latex(index=False)
180
+ if filename:
181
+ with open(filename, "w") as f:
182
+ f.write(latex_output)
183
+ return None
184
+ return latex_output
185
+
186
+ if csv:
187
+ if filename:
188
+ df.to_csv(filename, index=False)
189
+ return None
190
+
191
+ return df.to_csv(index=False)
192
+
193
+ return df
194
+
195
+ def show_schema(
196
+ self, shape: Literal["wide", "long"], remove_prefix: bool = False
197
+ ) -> None:
198
+ """Show the schema of the Results database.
199
+
200
+ :param shape: The shape of the data in the database (wide or long)
201
+ :param remove_prefix: Whether to remove the prefix from the column names
202
+
203
+ >>> from edsl.results import Results
204
+ >>> r = Results.example()
205
+ >>> r.show_schema(shape="long")
206
+ Type: table, Name: self, SQL: CREATE TABLE self (
207
+ ...
208
+ <BLANKLINE>
209
+ """
210
+ import pandas as pd
211
+
212
+ shape_enum = self._get_shape_enum(shape)
213
+ conn = self._db(shape=shape_enum, remove_prefix=remove_prefix)
214
+
215
+ if shape_enum == SQLDataShape.LONG:
216
+ # Query to get the schema of all tables
217
+ query = "SELECT type, name, sql FROM sqlite_master WHERE type='table'"
218
+ cursor = conn.execute(query)
219
+ schema = cursor.fetchall()
220
+ conn.close()
221
+
222
+ # Format and return the schema information
223
+ schema_info = ""
224
+ for row in schema:
225
+ schema_info += f"Type: {row[0]}, Name: {row[1]}, SQL: {row[2]}\n"
226
+
227
+ print(schema_info)
228
+ elif shape_enum == SQLDataShape.WIDE:
229
+ query = f"PRAGMA table_info(self)"
230
+ schema = pd.read_sql(query, conn)
231
+ # print(schema)
232
+ return schema
233
+
234
+
235
+ if __name__ == "__main__":
236
+ import doctest
237
+
238
+ doctest.testmod(optionflags=doctest.ELLIPSIS)
@@ -1,43 +1,43 @@
1
- """Mixin class for exporting results."""
2
-
3
- from functools import wraps
4
- from typing import Literal, Optional, Union
5
-
6
- from edsl.results.DatasetExportMixin import DatasetExportMixin
7
-
8
-
9
- def to_dataset(func):
10
- """Convert the Results object to a Dataset object before calling the function."""
11
-
12
- @wraps(func)
13
- def wrapper(self, *args, **kwargs):
14
- """Return the function with the Results object converted to a Dataset object."""
15
- if self.__class__.__name__ == "Results":
16
- return func(self.select(), *args, **kwargs)
17
- else:
18
- return func(self, *args, **kwargs)
19
-
20
- wrapper._is_wrapped = True
21
- return wrapper
22
-
23
-
24
- def decorate_methods_from_mixin(cls, mixin_cls):
25
- for attr_name, attr_value in mixin_cls.__dict__.items():
26
- if callable(attr_value) and not attr_name.startswith("__"):
27
- setattr(cls, attr_name, to_dataset(attr_value))
28
- return cls
29
-
30
-
31
- class ResultsExportMixin(DatasetExportMixin):
32
- """Mixin class for exporting Results objects."""
33
-
34
- def __init_subclass__(cls, **kwargs):
35
- super().__init_subclass__(**kwargs)
36
- decorate_methods_from_mixin(cls, DatasetExportMixin)
37
-
38
-
39
- if __name__ == "__main__":
40
- # pass
41
- import doctest
42
-
43
- doctest.testmod(optionflags=doctest.ELLIPSIS)
1
+ """Mixin class for exporting results."""
2
+
3
+ from functools import wraps
4
+ from typing import Literal, Optional, Union
5
+
6
+ from edsl.results.DatasetExportMixin import DatasetExportMixin
7
+
8
+
9
+ def to_dataset(func):
10
+ """Convert the Results object to a Dataset object before calling the function."""
11
+
12
+ @wraps(func)
13
+ def wrapper(self, *args, **kwargs):
14
+ """Return the function with the Results object converted to a Dataset object."""
15
+ if self.__class__.__name__ == "Results":
16
+ return func(self.select(), *args, **kwargs)
17
+ else:
18
+ return func(self, *args, **kwargs)
19
+
20
+ wrapper._is_wrapped = True
21
+ return wrapper
22
+
23
+
24
+ def decorate_methods_from_mixin(cls, mixin_cls):
25
+ for attr_name, attr_value in mixin_cls.__dict__.items():
26
+ if callable(attr_value) and not attr_name.startswith("__"):
27
+ setattr(cls, attr_name, to_dataset(attr_value))
28
+ return cls
29
+
30
+
31
+ class ResultsExportMixin(DatasetExportMixin):
32
+ """Mixin class for exporting Results objects."""
33
+
34
+ def __init_subclass__(cls, **kwargs):
35
+ super().__init_subclass__(**kwargs)
36
+ decorate_methods_from_mixin(cls, DatasetExportMixin)
37
+
38
+
39
+ if __name__ == "__main__":
40
+ # pass
41
+ import doctest
42
+
43
+ doctest.testmod(optionflags=doctest.ELLIPSIS)
@@ -1,33 +1,33 @@
1
- """Mixin for fetching data from results."""
2
-
3
- from functools import partial
4
- from itertools import chain
5
-
6
-
7
- class ResultsFetchMixin:
8
- """Mixin for fetching data from results."""
9
-
10
- def _fetch_list(self, data_type: str, key: str) -> list:
11
- """
12
- Return a list of values from the data for a given data type and key.
13
-
14
- Uses the filtered data, not the original data.
15
-
16
- Example:
17
-
18
- >>> from edsl.results import Results
19
- >>> r = Results.example()
20
- >>> r._fetch_list('answer', 'how_feeling')
21
- ['OK', 'Great', 'Terrible', 'OK']
22
- """
23
- returned_list = []
24
- for row in self.data:
25
- returned_list.append(row.sub_dicts[data_type].get(key, None))
26
-
27
- return returned_list
28
-
29
-
30
- if __name__ == "__main__":
31
- import doctest
32
-
33
- doctest.testmod(optionflags=doctest.ELLIPSIS)
1
+ """Mixin for fetching data from results."""
2
+
3
+ from functools import partial
4
+ from itertools import chain
5
+
6
+
7
+ class ResultsFetchMixin:
8
+ """Mixin for fetching data from results."""
9
+
10
+ def _fetch_list(self, data_type: str, key: str) -> list:
11
+ """
12
+ Return a list of values from the data for a given data type and key.
13
+
14
+ Uses the filtered data, not the original data.
15
+
16
+ Example:
17
+
18
+ >>> from edsl.results import Results
19
+ >>> r = Results.example()
20
+ >>> r._fetch_list('answer', 'how_feeling')
21
+ ['OK', 'Great', 'Terrible', 'OK']
22
+ """
23
+ returned_list = []
24
+ for row in self.data:
25
+ returned_list.append(row.sub_dicts[data_type].get(key, None))
26
+
27
+ return returned_list
28
+
29
+
30
+ if __name__ == "__main__":
31
+ import doctest
32
+
33
+ doctest.testmod(optionflags=doctest.ELLIPSIS)