edsl 0.1.37.dev5__py3-none-any.whl → 0.1.37.dev6__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 (261) hide show
  1. edsl/Base.py +303 -303
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +48 -48
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +855 -855
  7. edsl/agents/AgentList.py +350 -350
  8. edsl/agents/Invigilator.py +222 -222
  9. edsl/agents/InvigilatorBase.py +284 -284
  10. edsl/agents/PromptConstructor.py +353 -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 +289 -289
  26. edsl/config.py +149 -149
  27. edsl/conjure/AgentConstructionMixin.py +160 -160
  28. edsl/conjure/Conjure.py +62 -62
  29. edsl/conjure/InputData.py +659 -659
  30. edsl/conjure/InputDataCSV.py +48 -48
  31. edsl/conjure/InputDataMixinQuestionStats.py +182 -182
  32. edsl/conjure/InputDataPyRead.py +91 -91
  33. edsl/conjure/InputDataSPSS.py +8 -8
  34. edsl/conjure/InputDataStata.py +8 -8
  35. edsl/conjure/QuestionOptionMixin.py +76 -76
  36. edsl/conjure/QuestionTypeMixin.py +23 -23
  37. edsl/conjure/RawQuestion.py +65 -65
  38. edsl/conjure/SurveyResponses.py +7 -7
  39. edsl/conjure/__init__.py +9 -9
  40. edsl/conjure/naming_utilities.py +263 -263
  41. edsl/conjure/utilities.py +201 -201
  42. edsl/conversation/Conversation.py +290 -290
  43. edsl/conversation/car_buying.py +58 -58
  44. edsl/conversation/chips.py +95 -95
  45. edsl/conversation/mug_negotiation.py +81 -81
  46. edsl/conversation/next_speaker_utilities.py +93 -93
  47. edsl/coop/PriceFetcher.py +54 -54
  48. edsl/coop/__init__.py +2 -2
  49. edsl/coop/coop.py +958 -958
  50. edsl/coop/utils.py +131 -131
  51. edsl/data/Cache.py +527 -527
  52. edsl/data/CacheEntry.py +228 -228
  53. edsl/data/CacheHandler.py +149 -149
  54. edsl/data/RemoteCacheSync.py +97 -97
  55. edsl/data/SQLiteDict.py +292 -292
  56. edsl/data/__init__.py +4 -4
  57. edsl/data/orm.py +10 -10
  58. edsl/data_transfer_models.py +73 -73
  59. edsl/enums.py +173 -173
  60. edsl/exceptions/BaseException.py +21 -21
  61. edsl/exceptions/__init__.py +54 -54
  62. edsl/exceptions/agents.py +38 -38
  63. edsl/exceptions/configuration.py +16 -16
  64. edsl/exceptions/coop.py +10 -10
  65. edsl/exceptions/data.py +14 -14
  66. edsl/exceptions/general.py +34 -34
  67. edsl/exceptions/jobs.py +33 -33
  68. edsl/exceptions/language_models.py +63 -63
  69. edsl/exceptions/prompts.py +15 -15
  70. edsl/exceptions/questions.py +91 -91
  71. edsl/exceptions/results.py +29 -29
  72. edsl/exceptions/scenarios.py +22 -22
  73. edsl/exceptions/surveys.py +37 -37
  74. edsl/inference_services/AnthropicService.py +87 -87
  75. edsl/inference_services/AwsBedrock.py +120 -120
  76. edsl/inference_services/AzureAI.py +217 -217
  77. edsl/inference_services/DeepInfraService.py +18 -18
  78. edsl/inference_services/GoogleService.py +156 -156
  79. edsl/inference_services/GroqService.py +20 -20
  80. edsl/inference_services/InferenceServiceABC.py +147 -147
  81. edsl/inference_services/InferenceServicesCollection.py +97 -97
  82. edsl/inference_services/MistralAIService.py +123 -123
  83. edsl/inference_services/OllamaService.py +18 -18
  84. edsl/inference_services/OpenAIService.py +224 -224
  85. edsl/inference_services/TestService.py +89 -89
  86. edsl/inference_services/TogetherAIService.py +170 -170
  87. edsl/inference_services/models_available_cache.py +118 -118
  88. edsl/inference_services/rate_limits_cache.py +25 -25
  89. edsl/inference_services/registry.py +39 -39
  90. edsl/inference_services/write_available.py +10 -10
  91. edsl/jobs/Answers.py +56 -56
  92. edsl/jobs/Jobs.py +1347 -1347
  93. edsl/jobs/__init__.py +1 -1
  94. edsl/jobs/buckets/BucketCollection.py +63 -63
  95. edsl/jobs/buckets/ModelBuckets.py +65 -65
  96. edsl/jobs/buckets/TokenBucket.py +248 -248
  97. edsl/jobs/interviews/Interview.py +661 -661
  98. edsl/jobs/interviews/InterviewExceptionCollection.py +99 -99
  99. edsl/jobs/interviews/InterviewExceptionEntry.py +186 -186
  100. edsl/jobs/interviews/InterviewStatistic.py +63 -63
  101. edsl/jobs/interviews/InterviewStatisticsCollection.py +25 -25
  102. edsl/jobs/interviews/InterviewStatusDictionary.py +78 -78
  103. edsl/jobs/interviews/InterviewStatusLog.py +92 -92
  104. edsl/jobs/interviews/ReportErrors.py +66 -66
  105. edsl/jobs/interviews/interview_status_enum.py +9 -9
  106. edsl/jobs/runners/JobsRunnerAsyncio.py +338 -338
  107. edsl/jobs/runners/JobsRunnerStatus.py +332 -332
  108. edsl/jobs/tasks/QuestionTaskCreator.py +242 -242
  109. edsl/jobs/tasks/TaskCreators.py +64 -64
  110. edsl/jobs/tasks/TaskHistory.py +442 -442
  111. edsl/jobs/tasks/TaskStatusLog.py +23 -23
  112. edsl/jobs/tasks/task_status_enum.py +163 -163
  113. edsl/jobs/tokens/InterviewTokenUsage.py +27 -27
  114. edsl/jobs/tokens/TokenUsage.py +34 -34
  115. edsl/language_models/KeyLookup.py +30 -30
  116. edsl/language_models/LanguageModel.py +706 -706
  117. edsl/language_models/ModelList.py +102 -102
  118. edsl/language_models/RegisterLanguageModelsMeta.py +184 -184
  119. edsl/language_models/__init__.py +3 -3
  120. edsl/language_models/fake_openai_call.py +15 -15
  121. edsl/language_models/fake_openai_service.py +61 -61
  122. edsl/language_models/registry.py +137 -137
  123. edsl/language_models/repair.py +156 -156
  124. edsl/language_models/unused/ReplicateBase.py +83 -83
  125. edsl/language_models/utilities.py +64 -64
  126. edsl/notebooks/Notebook.py +259 -259
  127. edsl/notebooks/__init__.py +1 -1
  128. edsl/prompts/Prompt.py +357 -357
  129. edsl/prompts/__init__.py +2 -2
  130. edsl/questions/AnswerValidatorMixin.py +289 -289
  131. edsl/questions/QuestionBase.py +656 -656
  132. edsl/questions/QuestionBaseGenMixin.py +161 -161
  133. edsl/questions/QuestionBasePromptsMixin.py +234 -234
  134. edsl/questions/QuestionBudget.py +227 -227
  135. edsl/questions/QuestionCheckBox.py +359 -359
  136. edsl/questions/QuestionExtract.py +183 -183
  137. edsl/questions/QuestionFreeText.py +114 -114
  138. edsl/questions/QuestionFunctional.py +159 -159
  139. edsl/questions/QuestionList.py +231 -231
  140. edsl/questions/QuestionMultipleChoice.py +286 -286
  141. edsl/questions/QuestionNumerical.py +153 -153
  142. edsl/questions/QuestionRank.py +324 -324
  143. edsl/questions/Quick.py +41 -41
  144. edsl/questions/RegisterQuestionsMeta.py +71 -71
  145. edsl/questions/ResponseValidatorABC.py +174 -174
  146. edsl/questions/SimpleAskMixin.py +73 -73
  147. edsl/questions/__init__.py +26 -26
  148. edsl/questions/compose_questions.py +98 -98
  149. edsl/questions/decorators.py +21 -21
  150. edsl/questions/derived/QuestionLikertFive.py +76 -76
  151. edsl/questions/derived/QuestionLinearScale.py +87 -87
  152. edsl/questions/derived/QuestionTopK.py +91 -91
  153. edsl/questions/derived/QuestionYesNo.py +82 -82
  154. edsl/questions/descriptors.py +413 -413
  155. edsl/questions/prompt_templates/question_budget.jinja +13 -13
  156. edsl/questions/prompt_templates/question_checkbox.jinja +32 -32
  157. edsl/questions/prompt_templates/question_extract.jinja +11 -11
  158. edsl/questions/prompt_templates/question_free_text.jinja +3 -3
  159. edsl/questions/prompt_templates/question_linear_scale.jinja +11 -11
  160. edsl/questions/prompt_templates/question_list.jinja +17 -17
  161. edsl/questions/prompt_templates/question_multiple_choice.jinja +33 -33
  162. edsl/questions/prompt_templates/question_numerical.jinja +36 -36
  163. edsl/questions/question_registry.py +147 -147
  164. edsl/questions/settings.py +12 -12
  165. edsl/questions/templates/budget/answering_instructions.jinja +7 -7
  166. edsl/questions/templates/budget/question_presentation.jinja +7 -7
  167. edsl/questions/templates/checkbox/answering_instructions.jinja +10 -10
  168. edsl/questions/templates/checkbox/question_presentation.jinja +22 -22
  169. edsl/questions/templates/extract/answering_instructions.jinja +7 -7
  170. edsl/questions/templates/likert_five/answering_instructions.jinja +10 -10
  171. edsl/questions/templates/likert_five/question_presentation.jinja +11 -11
  172. edsl/questions/templates/linear_scale/answering_instructions.jinja +5 -5
  173. edsl/questions/templates/linear_scale/question_presentation.jinja +5 -5
  174. edsl/questions/templates/list/answering_instructions.jinja +3 -3
  175. edsl/questions/templates/list/question_presentation.jinja +5 -5
  176. edsl/questions/templates/multiple_choice/answering_instructions.jinja +9 -9
  177. edsl/questions/templates/multiple_choice/question_presentation.jinja +11 -11
  178. edsl/questions/templates/numerical/answering_instructions.jinja +6 -6
  179. edsl/questions/templates/numerical/question_presentation.jinja +6 -6
  180. edsl/questions/templates/rank/answering_instructions.jinja +11 -11
  181. edsl/questions/templates/rank/question_presentation.jinja +15 -15
  182. edsl/questions/templates/top_k/answering_instructions.jinja +8 -8
  183. edsl/questions/templates/top_k/question_presentation.jinja +22 -22
  184. edsl/questions/templates/yes_no/answering_instructions.jinja +6 -6
  185. edsl/questions/templates/yes_no/question_presentation.jinja +11 -11
  186. edsl/results/Dataset.py +293 -293
  187. edsl/results/DatasetExportMixin.py +717 -717
  188. edsl/results/DatasetTree.py +145 -145
  189. edsl/results/Result.py +450 -450
  190. edsl/results/Results.py +1071 -1071
  191. edsl/results/ResultsDBMixin.py +238 -238
  192. edsl/results/ResultsExportMixin.py +43 -43
  193. edsl/results/ResultsFetchMixin.py +33 -33
  194. edsl/results/ResultsGGMixin.py +121 -121
  195. edsl/results/ResultsToolsMixin.py +98 -98
  196. edsl/results/Selector.py +135 -135
  197. edsl/results/__init__.py +2 -2
  198. edsl/results/tree_explore.py +115 -115
  199. edsl/scenarios/FileStore.py +458 -458
  200. edsl/scenarios/Scenario.py +546 -546
  201. edsl/scenarios/ScenarioHtmlMixin.py +64 -64
  202. edsl/scenarios/ScenarioList.py +1112 -1112
  203. edsl/scenarios/ScenarioListExportMixin.py +52 -52
  204. edsl/scenarios/ScenarioListPdfMixin.py +261 -261
  205. edsl/scenarios/__init__.py +4 -4
  206. edsl/shared.py +1 -1
  207. edsl/study/ObjectEntry.py +173 -173
  208. edsl/study/ProofOfWork.py +113 -113
  209. edsl/study/SnapShot.py +80 -80
  210. edsl/study/Study.py +528 -528
  211. edsl/study/__init__.py +4 -4
  212. edsl/surveys/DAG.py +148 -148
  213. edsl/surveys/Memory.py +31 -31
  214. edsl/surveys/MemoryPlan.py +244 -244
  215. edsl/surveys/Rule.py +330 -330
  216. edsl/surveys/RuleCollection.py +387 -387
  217. edsl/surveys/Survey.py +1795 -1795
  218. edsl/surveys/SurveyCSS.py +261 -261
  219. edsl/surveys/SurveyExportMixin.py +259 -259
  220. edsl/surveys/SurveyFlowVisualizationMixin.py +121 -121
  221. edsl/surveys/SurveyQualtricsImport.py +284 -284
  222. edsl/surveys/__init__.py +3 -3
  223. edsl/surveys/base.py +53 -53
  224. edsl/surveys/descriptors.py +56 -56
  225. edsl/surveys/instructions/ChangeInstruction.py +47 -47
  226. edsl/surveys/instructions/Instruction.py +51 -51
  227. edsl/surveys/instructions/InstructionCollection.py +77 -77
  228. edsl/templates/error_reporting/base.html +23 -23
  229. edsl/templates/error_reporting/exceptions_by_model.html +34 -34
  230. edsl/templates/error_reporting/exceptions_by_question_name.html +16 -16
  231. edsl/templates/error_reporting/exceptions_by_type.html +16 -16
  232. edsl/templates/error_reporting/interview_details.html +115 -115
  233. edsl/templates/error_reporting/interviews.html +9 -9
  234. edsl/templates/error_reporting/overview.html +4 -4
  235. edsl/templates/error_reporting/performance_plot.html +1 -1
  236. edsl/templates/error_reporting/report.css +73 -73
  237. edsl/templates/error_reporting/report.html +117 -117
  238. edsl/templates/error_reporting/report.js +25 -25
  239. edsl/tools/__init__.py +1 -1
  240. edsl/tools/clusters.py +192 -192
  241. edsl/tools/embeddings.py +27 -27
  242. edsl/tools/embeddings_plotting.py +118 -118
  243. edsl/tools/plotting.py +112 -112
  244. edsl/tools/summarize.py +18 -18
  245. edsl/utilities/SystemInfo.py +28 -28
  246. edsl/utilities/__init__.py +22 -22
  247. edsl/utilities/ast_utilities.py +25 -25
  248. edsl/utilities/data/Registry.py +6 -6
  249. edsl/utilities/data/__init__.py +1 -1
  250. edsl/utilities/data/scooter_results.json +1 -1
  251. edsl/utilities/decorators.py +77 -77
  252. edsl/utilities/gcp_bucket/cloud_storage.py +96 -96
  253. edsl/utilities/interface.py +627 -627
  254. edsl/utilities/repair_functions.py +28 -28
  255. edsl/utilities/restricted_python.py +70 -70
  256. edsl/utilities/utilities.py +409 -409
  257. {edsl-0.1.37.dev5.dist-info → edsl-0.1.37.dev6.dist-info}/LICENSE +21 -21
  258. {edsl-0.1.37.dev5.dist-info → edsl-0.1.37.dev6.dist-info}/METADATA +1 -1
  259. edsl-0.1.37.dev6.dist-info/RECORD +283 -0
  260. edsl-0.1.37.dev5.dist-info/RECORD +0 -283
  261. {edsl-0.1.37.dev5.dist-info → edsl-0.1.37.dev6.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)