relationalai 0.13.5__py3-none-any.whl → 1.0.0a1__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 (856) hide show
  1. relationalai/__init__.py +1 -256
  2. relationalai/config/__init__.py +56 -0
  3. relationalai/config/config.py +289 -0
  4. relationalai/config/config_fields.py +86 -0
  5. relationalai/config/connections/__init__.py +46 -0
  6. relationalai/config/connections/base.py +23 -0
  7. relationalai/config/connections/duckdb.py +29 -0
  8. relationalai/config/connections/snowflake.py +243 -0
  9. relationalai/config/external/__init__.py +17 -0
  10. relationalai/config/external/dbt_converter.py +101 -0
  11. relationalai/config/external/dbt_models.py +93 -0
  12. relationalai/config/external/snowflake_converter.py +41 -0
  13. relationalai/config/external/snowflake_models.py +85 -0
  14. relationalai/config/external/utils.py +19 -0
  15. relationalai/config/shims.py +1 -0
  16. relationalai/semantics/__init__.py +146 -22
  17. relationalai/semantics/backends/lqp/annotations.py +11 -0
  18. relationalai/semantics/backends/sql/sql_compiler.py +327 -0
  19. relationalai/semantics/frontend/base.py +1716 -0
  20. relationalai/semantics/frontend/core.py +179 -0
  21. relationalai/semantics/frontend/front_compiler.py +1313 -0
  22. relationalai/semantics/frontend/pprint.py +408 -0
  23. relationalai/semantics/metamodel/__init__.py +6 -40
  24. relationalai/semantics/metamodel/builtins.py +205 -772
  25. relationalai/semantics/metamodel/metamodel.py +437 -0
  26. relationalai/semantics/metamodel/metamodel_analyzer.py +519 -0
  27. relationalai/semantics/metamodel/pprint.py +412 -0
  28. relationalai/semantics/metamodel/rewriter.py +266 -0
  29. relationalai/semantics/metamodel/typer.py +1186 -0
  30. relationalai/semantics/std/__init__.py +60 -40
  31. relationalai/semantics/std/aggregates.py +149 -0
  32. relationalai/semantics/std/common.py +44 -0
  33. relationalai/semantics/std/constraints.py +37 -43
  34. relationalai/semantics/std/datetime.py +246 -135
  35. relationalai/semantics/std/decimals.py +45 -52
  36. relationalai/semantics/std/floats.py +13 -5
  37. relationalai/semantics/std/integers.py +26 -11
  38. relationalai/semantics/std/math.py +183 -112
  39. relationalai/semantics/std/numbers.py +86 -0
  40. relationalai/semantics/std/re.py +80 -62
  41. relationalai/semantics/std/strings.py +101 -46
  42. relationalai/shims/executor.py +161 -0
  43. relationalai/shims/helpers.py +126 -0
  44. relationalai/shims/hoister.py +221 -0
  45. relationalai/shims/mm2v0.py +1324 -0
  46. relationalai/tools/cli/__init__.py +6 -0
  47. relationalai/tools/cli/cli.py +90 -0
  48. relationalai/tools/cli/components/__init__.py +5 -0
  49. relationalai/tools/cli/components/progress_reader.py +1524 -0
  50. relationalai/tools/cli/components/utils.py +58 -0
  51. relationalai/tools/cli/config_template.py +45 -0
  52. relationalai/tools/cli/dev.py +19 -0
  53. relationalai/tools/debugger.py +289 -183
  54. relationalai/tools/typer_debugger.py +93 -0
  55. relationalai/util/dataclasses.py +43 -0
  56. relationalai/util/docutils.py +40 -0
  57. relationalai/util/error.py +199 -0
  58. relationalai/util/format.py +48 -109
  59. relationalai/util/naming.py +145 -0
  60. relationalai/util/python.py +35 -0
  61. relationalai/util/runtime.py +156 -0
  62. relationalai/util/schema.py +197 -0
  63. relationalai/util/source.py +185 -0
  64. relationalai/util/structures.py +163 -0
  65. relationalai/util/tracing.py +261 -0
  66. relationalai-1.0.0a1.dist-info/METADATA +44 -0
  67. relationalai-1.0.0a1.dist-info/RECORD +489 -0
  68. relationalai-1.0.0a1.dist-info/WHEEL +5 -0
  69. relationalai-1.0.0a1.dist-info/entry_points.txt +3 -0
  70. relationalai-1.0.0a1.dist-info/top_level.txt +2 -0
  71. v0/relationalai/__init__.py +216 -0
  72. v0/relationalai/clients/__init__.py +5 -0
  73. v0/relationalai/clients/azure.py +477 -0
  74. v0/relationalai/clients/client.py +912 -0
  75. v0/relationalai/clients/config.py +673 -0
  76. v0/relationalai/clients/direct_access_client.py +118 -0
  77. v0/relationalai/clients/hash_util.py +31 -0
  78. v0/relationalai/clients/local.py +571 -0
  79. v0/relationalai/clients/profile_polling.py +73 -0
  80. v0/relationalai/clients/result_helpers.py +420 -0
  81. v0/relationalai/clients/snowflake.py +3869 -0
  82. v0/relationalai/clients/types.py +113 -0
  83. v0/relationalai/clients/use_index_poller.py +980 -0
  84. v0/relationalai/clients/util.py +356 -0
  85. v0/relationalai/debugging.py +389 -0
  86. v0/relationalai/dsl.py +1749 -0
  87. v0/relationalai/early_access/builder/__init__.py +30 -0
  88. v0/relationalai/early_access/builder/builder/__init__.py +35 -0
  89. v0/relationalai/early_access/builder/snowflake/__init__.py +12 -0
  90. v0/relationalai/early_access/builder/std/__init__.py +25 -0
  91. v0/relationalai/early_access/builder/std/decimals/__init__.py +12 -0
  92. v0/relationalai/early_access/builder/std/integers/__init__.py +12 -0
  93. v0/relationalai/early_access/builder/std/math/__init__.py +12 -0
  94. v0/relationalai/early_access/builder/std/strings/__init__.py +14 -0
  95. v0/relationalai/early_access/devtools/__init__.py +12 -0
  96. v0/relationalai/early_access/devtools/benchmark_lqp/__init__.py +12 -0
  97. v0/relationalai/early_access/devtools/extract_lqp/__init__.py +12 -0
  98. v0/relationalai/early_access/dsl/adapters/orm/adapter_qb.py +427 -0
  99. v0/relationalai/early_access/dsl/adapters/orm/parser.py +636 -0
  100. v0/relationalai/early_access/dsl/adapters/owl/adapter.py +176 -0
  101. v0/relationalai/early_access/dsl/adapters/owl/parser.py +160 -0
  102. v0/relationalai/early_access/dsl/bindings/common.py +402 -0
  103. v0/relationalai/early_access/dsl/bindings/csv.py +170 -0
  104. v0/relationalai/early_access/dsl/bindings/legacy/binding_models.py +143 -0
  105. v0/relationalai/early_access/dsl/bindings/snowflake.py +64 -0
  106. v0/relationalai/early_access/dsl/codegen/binder.py +411 -0
  107. v0/relationalai/early_access/dsl/codegen/common.py +79 -0
  108. v0/relationalai/early_access/dsl/codegen/helpers.py +23 -0
  109. v0/relationalai/early_access/dsl/codegen/relations.py +700 -0
  110. v0/relationalai/early_access/dsl/codegen/weaver.py +417 -0
  111. v0/relationalai/early_access/dsl/core/builders/__init__.py +47 -0
  112. v0/relationalai/early_access/dsl/core/builders/logic.py +19 -0
  113. v0/relationalai/early_access/dsl/core/builders/scalar_constraint.py +11 -0
  114. v0/relationalai/early_access/dsl/core/constraints/predicate/atomic.py +455 -0
  115. v0/relationalai/early_access/dsl/core/constraints/predicate/universal.py +73 -0
  116. v0/relationalai/early_access/dsl/core/constraints/scalar.py +310 -0
  117. v0/relationalai/early_access/dsl/core/context.py +13 -0
  118. v0/relationalai/early_access/dsl/core/cset.py +132 -0
  119. v0/relationalai/early_access/dsl/core/exprs/__init__.py +116 -0
  120. v0/relationalai/early_access/dsl/core/exprs/relational.py +18 -0
  121. v0/relationalai/early_access/dsl/core/exprs/scalar.py +412 -0
  122. v0/relationalai/early_access/dsl/core/instances.py +44 -0
  123. v0/relationalai/early_access/dsl/core/logic/__init__.py +193 -0
  124. v0/relationalai/early_access/dsl/core/logic/aggregation.py +98 -0
  125. v0/relationalai/early_access/dsl/core/logic/exists.py +223 -0
  126. v0/relationalai/early_access/dsl/core/logic/helper.py +163 -0
  127. v0/relationalai/early_access/dsl/core/namespaces.py +32 -0
  128. v0/relationalai/early_access/dsl/core/relations.py +276 -0
  129. v0/relationalai/early_access/dsl/core/rules.py +112 -0
  130. v0/relationalai/early_access/dsl/core/std/__init__.py +45 -0
  131. v0/relationalai/early_access/dsl/core/temporal/recall.py +6 -0
  132. v0/relationalai/early_access/dsl/core/types/__init__.py +270 -0
  133. v0/relationalai/early_access/dsl/core/types/concepts.py +128 -0
  134. v0/relationalai/early_access/dsl/core/types/constrained/__init__.py +267 -0
  135. v0/relationalai/early_access/dsl/core/types/constrained/nominal.py +143 -0
  136. v0/relationalai/early_access/dsl/core/types/constrained/subtype.py +124 -0
  137. v0/relationalai/early_access/dsl/core/types/standard.py +92 -0
  138. v0/relationalai/early_access/dsl/core/types/unconstrained.py +50 -0
  139. v0/relationalai/early_access/dsl/core/types/variables.py +203 -0
  140. v0/relationalai/early_access/dsl/ir/compiler.py +318 -0
  141. v0/relationalai/early_access/dsl/ir/executor.py +260 -0
  142. v0/relationalai/early_access/dsl/ontologies/constraints.py +88 -0
  143. v0/relationalai/early_access/dsl/ontologies/export.py +30 -0
  144. v0/relationalai/early_access/dsl/ontologies/models.py +453 -0
  145. v0/relationalai/early_access/dsl/ontologies/python_printer.py +303 -0
  146. v0/relationalai/early_access/dsl/ontologies/readings.py +60 -0
  147. v0/relationalai/early_access/dsl/ontologies/relationships.py +322 -0
  148. v0/relationalai/early_access/dsl/ontologies/roles.py +87 -0
  149. v0/relationalai/early_access/dsl/ontologies/subtyping.py +55 -0
  150. v0/relationalai/early_access/dsl/orm/constraints.py +438 -0
  151. v0/relationalai/early_access/dsl/orm/measures/dimensions.py +200 -0
  152. v0/relationalai/early_access/dsl/orm/measures/initializer.py +16 -0
  153. v0/relationalai/early_access/dsl/orm/measures/measure_rules.py +275 -0
  154. v0/relationalai/early_access/dsl/orm/measures/measures.py +299 -0
  155. v0/relationalai/early_access/dsl/orm/measures/role_exprs.py +268 -0
  156. v0/relationalai/early_access/dsl/orm/models.py +256 -0
  157. v0/relationalai/early_access/dsl/orm/object_oriented_printer.py +344 -0
  158. v0/relationalai/early_access/dsl/orm/printer.py +469 -0
  159. v0/relationalai/early_access/dsl/orm/reasoners.py +480 -0
  160. v0/relationalai/early_access/dsl/orm/relations.py +19 -0
  161. v0/relationalai/early_access/dsl/orm/relationships.py +251 -0
  162. v0/relationalai/early_access/dsl/orm/types.py +42 -0
  163. v0/relationalai/early_access/dsl/orm/utils.py +79 -0
  164. v0/relationalai/early_access/dsl/orm/verb.py +204 -0
  165. v0/relationalai/early_access/dsl/physical_metadata/tables.py +133 -0
  166. v0/relationalai/early_access/dsl/relations.py +170 -0
  167. v0/relationalai/early_access/dsl/rulesets.py +69 -0
  168. v0/relationalai/early_access/dsl/schemas/__init__.py +450 -0
  169. v0/relationalai/early_access/dsl/schemas/builder.py +48 -0
  170. v0/relationalai/early_access/dsl/schemas/comp_names.py +51 -0
  171. v0/relationalai/early_access/dsl/schemas/components.py +203 -0
  172. v0/relationalai/early_access/dsl/schemas/contexts.py +156 -0
  173. v0/relationalai/early_access/dsl/schemas/exprs.py +89 -0
  174. v0/relationalai/early_access/dsl/schemas/fragments.py +464 -0
  175. v0/relationalai/early_access/dsl/serialization.py +79 -0
  176. v0/relationalai/early_access/dsl/serialize/exporter.py +163 -0
  177. v0/relationalai/early_access/dsl/snow/api.py +104 -0
  178. v0/relationalai/early_access/dsl/snow/common.py +76 -0
  179. v0/relationalai/early_access/dsl/state_mgmt/__init__.py +129 -0
  180. v0/relationalai/early_access/dsl/state_mgmt/state_charts.py +125 -0
  181. v0/relationalai/early_access/dsl/state_mgmt/transitions.py +130 -0
  182. v0/relationalai/early_access/dsl/types/__init__.py +40 -0
  183. v0/relationalai/early_access/dsl/types/concepts.py +12 -0
  184. v0/relationalai/early_access/dsl/types/entities.py +135 -0
  185. v0/relationalai/early_access/dsl/types/values.py +17 -0
  186. v0/relationalai/early_access/dsl/utils.py +102 -0
  187. v0/relationalai/early_access/graphs/__init__.py +13 -0
  188. v0/relationalai/early_access/lqp/__init__.py +12 -0
  189. v0/relationalai/early_access/lqp/compiler/__init__.py +12 -0
  190. v0/relationalai/early_access/lqp/constructors/__init__.py +18 -0
  191. v0/relationalai/early_access/lqp/executor/__init__.py +12 -0
  192. v0/relationalai/early_access/lqp/ir/__init__.py +12 -0
  193. v0/relationalai/early_access/lqp/passes/__init__.py +12 -0
  194. v0/relationalai/early_access/lqp/pragmas/__init__.py +12 -0
  195. v0/relationalai/early_access/lqp/primitives/__init__.py +12 -0
  196. v0/relationalai/early_access/lqp/types/__init__.py +12 -0
  197. v0/relationalai/early_access/lqp/utils/__init__.py +12 -0
  198. v0/relationalai/early_access/lqp/validators/__init__.py +12 -0
  199. v0/relationalai/early_access/metamodel/__init__.py +58 -0
  200. v0/relationalai/early_access/metamodel/builtins/__init__.py +12 -0
  201. v0/relationalai/early_access/metamodel/compiler/__init__.py +12 -0
  202. v0/relationalai/early_access/metamodel/dependency/__init__.py +12 -0
  203. v0/relationalai/early_access/metamodel/factory/__init__.py +17 -0
  204. v0/relationalai/early_access/metamodel/helpers/__init__.py +12 -0
  205. v0/relationalai/early_access/metamodel/ir/__init__.py +14 -0
  206. v0/relationalai/early_access/metamodel/rewrite/__init__.py +7 -0
  207. v0/relationalai/early_access/metamodel/typer/__init__.py +3 -0
  208. v0/relationalai/early_access/metamodel/typer/typer/__init__.py +12 -0
  209. v0/relationalai/early_access/metamodel/types/__init__.py +15 -0
  210. v0/relationalai/early_access/metamodel/util/__init__.py +15 -0
  211. v0/relationalai/early_access/metamodel/visitor/__init__.py +12 -0
  212. v0/relationalai/early_access/rel/__init__.py +12 -0
  213. v0/relationalai/early_access/rel/executor/__init__.py +12 -0
  214. v0/relationalai/early_access/rel/rel_utils/__init__.py +12 -0
  215. v0/relationalai/early_access/rel/rewrite/__init__.py +7 -0
  216. v0/relationalai/early_access/solvers/__init__.py +19 -0
  217. v0/relationalai/early_access/sql/__init__.py +11 -0
  218. v0/relationalai/early_access/sql/executor/__init__.py +3 -0
  219. v0/relationalai/early_access/sql/rewrite/__init__.py +3 -0
  220. v0/relationalai/early_access/tests/logging/__init__.py +12 -0
  221. v0/relationalai/early_access/tests/test_snapshot_base/__init__.py +12 -0
  222. v0/relationalai/early_access/tests/utils/__init__.py +12 -0
  223. v0/relationalai/environments/__init__.py +35 -0
  224. v0/relationalai/environments/base.py +381 -0
  225. v0/relationalai/environments/colab.py +14 -0
  226. v0/relationalai/environments/generic.py +71 -0
  227. v0/relationalai/environments/ipython.py +68 -0
  228. v0/relationalai/environments/jupyter.py +9 -0
  229. v0/relationalai/environments/snowbook.py +169 -0
  230. v0/relationalai/errors.py +2455 -0
  231. v0/relationalai/experimental/SF.py +38 -0
  232. v0/relationalai/experimental/inspect.py +47 -0
  233. v0/relationalai/experimental/pathfinder/__init__.py +158 -0
  234. v0/relationalai/experimental/pathfinder/api.py +160 -0
  235. v0/relationalai/experimental/pathfinder/automaton.py +584 -0
  236. v0/relationalai/experimental/pathfinder/bridge.py +226 -0
  237. v0/relationalai/experimental/pathfinder/compiler.py +416 -0
  238. v0/relationalai/experimental/pathfinder/datalog.py +214 -0
  239. v0/relationalai/experimental/pathfinder/diagnostics.py +56 -0
  240. v0/relationalai/experimental/pathfinder/filter.py +236 -0
  241. v0/relationalai/experimental/pathfinder/glushkov.py +439 -0
  242. v0/relationalai/experimental/pathfinder/options.py +265 -0
  243. v0/relationalai/experimental/pathfinder/rpq.py +344 -0
  244. v0/relationalai/experimental/pathfinder/transition.py +200 -0
  245. v0/relationalai/experimental/pathfinder/utils.py +26 -0
  246. v0/relationalai/experimental/paths/api.py +143 -0
  247. v0/relationalai/experimental/paths/benchmarks/grid_graph.py +37 -0
  248. v0/relationalai/experimental/paths/examples/basic_example.py +40 -0
  249. v0/relationalai/experimental/paths/examples/minimal_engine_warmup.py +3 -0
  250. v0/relationalai/experimental/paths/examples/movie_example.py +77 -0
  251. v0/relationalai/experimental/paths/examples/paths_benchmark.py +115 -0
  252. v0/relationalai/experimental/paths/examples/paths_example.py +116 -0
  253. v0/relationalai/experimental/paths/examples/pattern_to_automaton.py +28 -0
  254. v0/relationalai/experimental/paths/find_paths_via_automaton.py +85 -0
  255. v0/relationalai/experimental/paths/graph.py +185 -0
  256. v0/relationalai/experimental/paths/path_algorithms/find_paths.py +280 -0
  257. v0/relationalai/experimental/paths/path_algorithms/one_sided_ball_repetition.py +26 -0
  258. v0/relationalai/experimental/paths/path_algorithms/one_sided_ball_upto.py +111 -0
  259. v0/relationalai/experimental/paths/path_algorithms/single.py +59 -0
  260. v0/relationalai/experimental/paths/path_algorithms/two_sided_balls_repetition.py +39 -0
  261. v0/relationalai/experimental/paths/path_algorithms/two_sided_balls_upto.py +103 -0
  262. v0/relationalai/experimental/paths/path_algorithms/usp-old.py +130 -0
  263. v0/relationalai/experimental/paths/path_algorithms/usp-tuple.py +183 -0
  264. v0/relationalai/experimental/paths/path_algorithms/usp.py +150 -0
  265. v0/relationalai/experimental/paths/product_graph.py +93 -0
  266. v0/relationalai/experimental/paths/rpq/automaton.py +584 -0
  267. v0/relationalai/experimental/paths/rpq/diagnostics.py +56 -0
  268. v0/relationalai/experimental/paths/rpq/rpq.py +378 -0
  269. v0/relationalai/experimental/paths/tests/tests_limit_sp_max_length.py +90 -0
  270. v0/relationalai/experimental/paths/tests/tests_limit_sp_multiple.py +119 -0
  271. v0/relationalai/experimental/paths/tests/tests_limit_sp_single.py +104 -0
  272. v0/relationalai/experimental/paths/tests/tests_limit_walks_multiple.py +113 -0
  273. v0/relationalai/experimental/paths/tests/tests_limit_walks_single.py +149 -0
  274. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_multiple.py +70 -0
  275. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_single.py +64 -0
  276. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_upto_multiple.py +115 -0
  277. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_upto_single.py +75 -0
  278. v0/relationalai/experimental/paths/tests/tests_single_paths.py +152 -0
  279. v0/relationalai/experimental/paths/tests/tests_single_walks.py +208 -0
  280. v0/relationalai/experimental/paths/tests/tests_single_walks_undirected.py +297 -0
  281. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_multiple.py +107 -0
  282. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_single.py +76 -0
  283. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_upto_multiple.py +76 -0
  284. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_upto_single.py +110 -0
  285. v0/relationalai/experimental/paths/tests/tests_usp_nsp_multiple.py +229 -0
  286. v0/relationalai/experimental/paths/tests/tests_usp_nsp_single.py +108 -0
  287. v0/relationalai/experimental/paths/tree_agg.py +168 -0
  288. v0/relationalai/experimental/paths/utilities/iterators.py +27 -0
  289. v0/relationalai/experimental/paths/utilities/prefix_sum.py +91 -0
  290. v0/relationalai/experimental/solvers.py +1087 -0
  291. v0/relationalai/loaders/csv.py +195 -0
  292. v0/relationalai/loaders/loader.py +177 -0
  293. v0/relationalai/loaders/types.py +23 -0
  294. v0/relationalai/rel_emitter.py +373 -0
  295. v0/relationalai/rel_utils.py +185 -0
  296. v0/relationalai/semantics/__init__.py +29 -0
  297. v0/relationalai/semantics/devtools/benchmark_lqp.py +536 -0
  298. v0/relationalai/semantics/devtools/compilation_manager.py +294 -0
  299. v0/relationalai/semantics/devtools/extract_lqp.py +110 -0
  300. v0/relationalai/semantics/internal/internal.py +3785 -0
  301. v0/relationalai/semantics/internal/snowflake.py +324 -0
  302. v0/relationalai/semantics/lqp/builtins.py +16 -0
  303. v0/relationalai/semantics/lqp/compiler.py +22 -0
  304. v0/relationalai/semantics/lqp/constructors.py +68 -0
  305. v0/relationalai/semantics/lqp/executor.py +469 -0
  306. v0/relationalai/semantics/lqp/intrinsics.py +24 -0
  307. v0/relationalai/semantics/lqp/ir.py +124 -0
  308. v0/relationalai/semantics/lqp/model2lqp.py +839 -0
  309. v0/relationalai/semantics/lqp/passes.py +680 -0
  310. v0/relationalai/semantics/lqp/primitives.py +252 -0
  311. v0/relationalai/semantics/lqp/result_helpers.py +202 -0
  312. v0/relationalai/semantics/lqp/rewrite/__init__.py +18 -0
  313. v0/relationalai/semantics/lqp/rewrite/annotate_constraints.py +57 -0
  314. v0/relationalai/semantics/lqp/rewrite/cdc.py +216 -0
  315. v0/relationalai/semantics/lqp/rewrite/extract_common.py +338 -0
  316. v0/relationalai/semantics/lqp/rewrite/extract_keys.py +449 -0
  317. v0/relationalai/semantics/lqp/rewrite/function_annotations.py +114 -0
  318. v0/relationalai/semantics/lqp/rewrite/functional_dependencies.py +314 -0
  319. v0/relationalai/semantics/lqp/rewrite/quantify_vars.py +296 -0
  320. v0/relationalai/semantics/lqp/rewrite/splinter.py +76 -0
  321. v0/relationalai/semantics/lqp/types.py +101 -0
  322. v0/relationalai/semantics/lqp/utils.py +160 -0
  323. v0/relationalai/semantics/lqp/validators.py +57 -0
  324. v0/relationalai/semantics/metamodel/__init__.py +40 -0
  325. v0/relationalai/semantics/metamodel/builtins.py +774 -0
  326. v0/relationalai/semantics/metamodel/compiler.py +133 -0
  327. v0/relationalai/semantics/metamodel/dependency.py +862 -0
  328. v0/relationalai/semantics/metamodel/executor.py +61 -0
  329. v0/relationalai/semantics/metamodel/factory.py +287 -0
  330. v0/relationalai/semantics/metamodel/helpers.py +361 -0
  331. v0/relationalai/semantics/metamodel/ir.py +923 -0
  332. v0/relationalai/semantics/metamodel/rewrite/__init__.py +7 -0
  333. v0/relationalai/semantics/metamodel/rewrite/discharge_constraints.py +39 -0
  334. v0/relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py +210 -0
  335. v0/relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py +78 -0
  336. v0/relationalai/semantics/metamodel/rewrite/flatten.py +549 -0
  337. v0/relationalai/semantics/metamodel/rewrite/format_outputs.py +165 -0
  338. v0/relationalai/semantics/metamodel/typer/checker.py +353 -0
  339. v0/relationalai/semantics/metamodel/typer/typer.py +1395 -0
  340. v0/relationalai/semantics/metamodel/util.py +505 -0
  341. v0/relationalai/semantics/metamodel/visitor.py +944 -0
  342. v0/relationalai/semantics/reasoners/__init__.py +10 -0
  343. v0/relationalai/semantics/reasoners/graph/__init__.py +37 -0
  344. v0/relationalai/semantics/reasoners/graph/core.py +9020 -0
  345. v0/relationalai/semantics/reasoners/optimization/__init__.py +68 -0
  346. v0/relationalai/semantics/reasoners/optimization/common.py +88 -0
  347. v0/relationalai/semantics/reasoners/optimization/solvers_dev.py +568 -0
  348. v0/relationalai/semantics/reasoners/optimization/solvers_pb.py +1163 -0
  349. v0/relationalai/semantics/rel/builtins.py +40 -0
  350. v0/relationalai/semantics/rel/compiler.py +989 -0
  351. v0/relationalai/semantics/rel/executor.py +359 -0
  352. v0/relationalai/semantics/rel/rel.py +482 -0
  353. v0/relationalai/semantics/rel/rel_utils.py +276 -0
  354. v0/relationalai/semantics/snowflake/__init__.py +3 -0
  355. v0/relationalai/semantics/sql/compiler.py +2503 -0
  356. v0/relationalai/semantics/sql/executor/duck_db.py +52 -0
  357. v0/relationalai/semantics/sql/executor/result_helpers.py +64 -0
  358. v0/relationalai/semantics/sql/executor/snowflake.py +145 -0
  359. v0/relationalai/semantics/sql/rewrite/denormalize.py +222 -0
  360. v0/relationalai/semantics/sql/rewrite/double_negation.py +49 -0
  361. v0/relationalai/semantics/sql/rewrite/recursive_union.py +127 -0
  362. v0/relationalai/semantics/sql/rewrite/sort_output_query.py +246 -0
  363. v0/relationalai/semantics/sql/sql.py +504 -0
  364. v0/relationalai/semantics/std/__init__.py +54 -0
  365. v0/relationalai/semantics/std/constraints.py +43 -0
  366. v0/relationalai/semantics/std/datetime.py +363 -0
  367. v0/relationalai/semantics/std/decimals.py +62 -0
  368. v0/relationalai/semantics/std/floats.py +7 -0
  369. v0/relationalai/semantics/std/integers.py +22 -0
  370. v0/relationalai/semantics/std/math.py +141 -0
  371. v0/relationalai/semantics/std/pragmas.py +11 -0
  372. v0/relationalai/semantics/std/re.py +83 -0
  373. v0/relationalai/semantics/std/std.py +14 -0
  374. v0/relationalai/semantics/std/strings.py +63 -0
  375. v0/relationalai/semantics/tests/__init__.py +0 -0
  376. v0/relationalai/semantics/tests/test_snapshot_abstract.py +143 -0
  377. v0/relationalai/semantics/tests/test_snapshot_base.py +9 -0
  378. v0/relationalai/semantics/tests/utils.py +46 -0
  379. v0/relationalai/std/__init__.py +70 -0
  380. v0/relationalai/tools/__init__.py +0 -0
  381. v0/relationalai/tools/cli.py +1940 -0
  382. v0/relationalai/tools/cli_controls.py +1826 -0
  383. v0/relationalai/tools/cli_helpers.py +390 -0
  384. v0/relationalai/tools/debugger.py +183 -0
  385. v0/relationalai/tools/debugger_client.py +109 -0
  386. v0/relationalai/tools/debugger_server.py +302 -0
  387. v0/relationalai/tools/dev.py +685 -0
  388. v0/relationalai/tools/qb_debugger.py +425 -0
  389. v0/relationalai/util/clean_up_databases.py +95 -0
  390. v0/relationalai/util/format.py +123 -0
  391. v0/relationalai/util/list_databases.py +9 -0
  392. v0/relationalai/util/otel_configuration.py +25 -0
  393. v0/relationalai/util/otel_handler.py +484 -0
  394. v0/relationalai/util/snowflake_handler.py +88 -0
  395. v0/relationalai/util/span_format_test.py +43 -0
  396. v0/relationalai/util/span_tracker.py +207 -0
  397. v0/relationalai/util/spans_file_handler.py +72 -0
  398. v0/relationalai/util/tracing_handler.py +34 -0
  399. frontend/debugger/dist/.gitignore +0 -2
  400. frontend/debugger/dist/assets/favicon-Dy0ZgA6N.png +0 -0
  401. frontend/debugger/dist/assets/index-Cssla-O7.js +0 -208
  402. frontend/debugger/dist/assets/index-DlHsYx1V.css +0 -9
  403. frontend/debugger/dist/index.html +0 -17
  404. relationalai/clients/__init__.py +0 -18
  405. relationalai/clients/client.py +0 -946
  406. relationalai/clients/config.py +0 -673
  407. relationalai/clients/direct_access_client.py +0 -118
  408. relationalai/clients/exec_txn_poller.py +0 -153
  409. relationalai/clients/hash_util.py +0 -31
  410. relationalai/clients/local.py +0 -594
  411. relationalai/clients/profile_polling.py +0 -73
  412. relationalai/clients/resources/__init__.py +0 -8
  413. relationalai/clients/resources/azure/azure.py +0 -502
  414. relationalai/clients/resources/snowflake/__init__.py +0 -20
  415. relationalai/clients/resources/snowflake/cli_resources.py +0 -98
  416. relationalai/clients/resources/snowflake/direct_access_resources.py +0 -739
  417. relationalai/clients/resources/snowflake/engine_service.py +0 -381
  418. relationalai/clients/resources/snowflake/engine_state_handlers.py +0 -315
  419. relationalai/clients/resources/snowflake/error_handlers.py +0 -240
  420. relationalai/clients/resources/snowflake/export_procedure.py.jinja +0 -249
  421. relationalai/clients/resources/snowflake/resources_factory.py +0 -99
  422. relationalai/clients/resources/snowflake/snowflake.py +0 -3193
  423. relationalai/clients/resources/snowflake/use_index_poller.py +0 -1019
  424. relationalai/clients/resources/snowflake/use_index_resources.py +0 -188
  425. relationalai/clients/resources/snowflake/util.py +0 -387
  426. relationalai/clients/result_helpers.py +0 -420
  427. relationalai/clients/types.py +0 -118
  428. relationalai/clients/util.py +0 -356
  429. relationalai/debugging.py +0 -389
  430. relationalai/dsl.py +0 -1749
  431. relationalai/early_access/builder/__init__.py +0 -30
  432. relationalai/early_access/builder/builder/__init__.py +0 -35
  433. relationalai/early_access/builder/snowflake/__init__.py +0 -12
  434. relationalai/early_access/builder/std/__init__.py +0 -25
  435. relationalai/early_access/builder/std/decimals/__init__.py +0 -12
  436. relationalai/early_access/builder/std/integers/__init__.py +0 -12
  437. relationalai/early_access/builder/std/math/__init__.py +0 -12
  438. relationalai/early_access/builder/std/strings/__init__.py +0 -14
  439. relationalai/early_access/devtools/__init__.py +0 -12
  440. relationalai/early_access/devtools/benchmark_lqp/__init__.py +0 -12
  441. relationalai/early_access/devtools/extract_lqp/__init__.py +0 -12
  442. relationalai/early_access/dsl/adapters/orm/adapter_qb.py +0 -427
  443. relationalai/early_access/dsl/adapters/orm/parser.py +0 -636
  444. relationalai/early_access/dsl/adapters/owl/adapter.py +0 -176
  445. relationalai/early_access/dsl/adapters/owl/parser.py +0 -160
  446. relationalai/early_access/dsl/bindings/common.py +0 -402
  447. relationalai/early_access/dsl/bindings/csv.py +0 -170
  448. relationalai/early_access/dsl/bindings/legacy/binding_models.py +0 -143
  449. relationalai/early_access/dsl/bindings/snowflake.py +0 -64
  450. relationalai/early_access/dsl/codegen/binder.py +0 -411
  451. relationalai/early_access/dsl/codegen/common.py +0 -79
  452. relationalai/early_access/dsl/codegen/helpers.py +0 -23
  453. relationalai/early_access/dsl/codegen/relations.py +0 -700
  454. relationalai/early_access/dsl/codegen/weaver.py +0 -417
  455. relationalai/early_access/dsl/core/builders/__init__.py +0 -47
  456. relationalai/early_access/dsl/core/builders/logic.py +0 -19
  457. relationalai/early_access/dsl/core/builders/scalar_constraint.py +0 -11
  458. relationalai/early_access/dsl/core/constraints/predicate/atomic.py +0 -455
  459. relationalai/early_access/dsl/core/constraints/predicate/universal.py +0 -73
  460. relationalai/early_access/dsl/core/constraints/scalar.py +0 -310
  461. relationalai/early_access/dsl/core/context.py +0 -13
  462. relationalai/early_access/dsl/core/cset.py +0 -132
  463. relationalai/early_access/dsl/core/exprs/__init__.py +0 -116
  464. relationalai/early_access/dsl/core/exprs/relational.py +0 -18
  465. relationalai/early_access/dsl/core/exprs/scalar.py +0 -412
  466. relationalai/early_access/dsl/core/instances.py +0 -44
  467. relationalai/early_access/dsl/core/logic/__init__.py +0 -193
  468. relationalai/early_access/dsl/core/logic/aggregation.py +0 -98
  469. relationalai/early_access/dsl/core/logic/exists.py +0 -223
  470. relationalai/early_access/dsl/core/logic/helper.py +0 -163
  471. relationalai/early_access/dsl/core/namespaces.py +0 -32
  472. relationalai/early_access/dsl/core/relations.py +0 -276
  473. relationalai/early_access/dsl/core/rules.py +0 -112
  474. relationalai/early_access/dsl/core/std/__init__.py +0 -45
  475. relationalai/early_access/dsl/core/temporal/recall.py +0 -6
  476. relationalai/early_access/dsl/core/types/__init__.py +0 -270
  477. relationalai/early_access/dsl/core/types/concepts.py +0 -128
  478. relationalai/early_access/dsl/core/types/constrained/__init__.py +0 -267
  479. relationalai/early_access/dsl/core/types/constrained/nominal.py +0 -143
  480. relationalai/early_access/dsl/core/types/constrained/subtype.py +0 -124
  481. relationalai/early_access/dsl/core/types/standard.py +0 -92
  482. relationalai/early_access/dsl/core/types/unconstrained.py +0 -50
  483. relationalai/early_access/dsl/core/types/variables.py +0 -203
  484. relationalai/early_access/dsl/ir/compiler.py +0 -318
  485. relationalai/early_access/dsl/ir/executor.py +0 -260
  486. relationalai/early_access/dsl/ontologies/constraints.py +0 -88
  487. relationalai/early_access/dsl/ontologies/export.py +0 -30
  488. relationalai/early_access/dsl/ontologies/models.py +0 -453
  489. relationalai/early_access/dsl/ontologies/python_printer.py +0 -303
  490. relationalai/early_access/dsl/ontologies/readings.py +0 -60
  491. relationalai/early_access/dsl/ontologies/relationships.py +0 -322
  492. relationalai/early_access/dsl/ontologies/roles.py +0 -87
  493. relationalai/early_access/dsl/ontologies/subtyping.py +0 -55
  494. relationalai/early_access/dsl/orm/constraints.py +0 -438
  495. relationalai/early_access/dsl/orm/measures/dimensions.py +0 -200
  496. relationalai/early_access/dsl/orm/measures/initializer.py +0 -16
  497. relationalai/early_access/dsl/orm/measures/measure_rules.py +0 -275
  498. relationalai/early_access/dsl/orm/measures/measures.py +0 -299
  499. relationalai/early_access/dsl/orm/measures/role_exprs.py +0 -268
  500. relationalai/early_access/dsl/orm/models.py +0 -256
  501. relationalai/early_access/dsl/orm/object_oriented_printer.py +0 -344
  502. relationalai/early_access/dsl/orm/printer.py +0 -469
  503. relationalai/early_access/dsl/orm/reasoners.py +0 -480
  504. relationalai/early_access/dsl/orm/relations.py +0 -19
  505. relationalai/early_access/dsl/orm/relationships.py +0 -251
  506. relationalai/early_access/dsl/orm/types.py +0 -42
  507. relationalai/early_access/dsl/orm/utils.py +0 -79
  508. relationalai/early_access/dsl/orm/verb.py +0 -204
  509. relationalai/early_access/dsl/physical_metadata/tables.py +0 -133
  510. relationalai/early_access/dsl/relations.py +0 -170
  511. relationalai/early_access/dsl/rulesets.py +0 -69
  512. relationalai/early_access/dsl/schemas/__init__.py +0 -450
  513. relationalai/early_access/dsl/schemas/builder.py +0 -48
  514. relationalai/early_access/dsl/schemas/comp_names.py +0 -51
  515. relationalai/early_access/dsl/schemas/components.py +0 -203
  516. relationalai/early_access/dsl/schemas/contexts.py +0 -156
  517. relationalai/early_access/dsl/schemas/exprs.py +0 -89
  518. relationalai/early_access/dsl/schemas/fragments.py +0 -464
  519. relationalai/early_access/dsl/serialization.py +0 -79
  520. relationalai/early_access/dsl/serialize/exporter.py +0 -163
  521. relationalai/early_access/dsl/snow/api.py +0 -105
  522. relationalai/early_access/dsl/snow/common.py +0 -76
  523. relationalai/early_access/dsl/state_mgmt/__init__.py +0 -129
  524. relationalai/early_access/dsl/state_mgmt/state_charts.py +0 -125
  525. relationalai/early_access/dsl/state_mgmt/transitions.py +0 -130
  526. relationalai/early_access/dsl/types/__init__.py +0 -40
  527. relationalai/early_access/dsl/types/concepts.py +0 -12
  528. relationalai/early_access/dsl/types/entities.py +0 -135
  529. relationalai/early_access/dsl/types/values.py +0 -17
  530. relationalai/early_access/dsl/utils.py +0 -102
  531. relationalai/early_access/graphs/__init__.py +0 -13
  532. relationalai/early_access/lqp/__init__.py +0 -12
  533. relationalai/early_access/lqp/compiler/__init__.py +0 -12
  534. relationalai/early_access/lqp/constructors/__init__.py +0 -18
  535. relationalai/early_access/lqp/executor/__init__.py +0 -12
  536. relationalai/early_access/lqp/ir/__init__.py +0 -12
  537. relationalai/early_access/lqp/passes/__init__.py +0 -12
  538. relationalai/early_access/lqp/pragmas/__init__.py +0 -12
  539. relationalai/early_access/lqp/primitives/__init__.py +0 -12
  540. relationalai/early_access/lqp/types/__init__.py +0 -12
  541. relationalai/early_access/lqp/utils/__init__.py +0 -12
  542. relationalai/early_access/lqp/validators/__init__.py +0 -12
  543. relationalai/early_access/metamodel/__init__.py +0 -58
  544. relationalai/early_access/metamodel/builtins/__init__.py +0 -12
  545. relationalai/early_access/metamodel/compiler/__init__.py +0 -12
  546. relationalai/early_access/metamodel/dependency/__init__.py +0 -12
  547. relationalai/early_access/metamodel/factory/__init__.py +0 -17
  548. relationalai/early_access/metamodel/helpers/__init__.py +0 -12
  549. relationalai/early_access/metamodel/ir/__init__.py +0 -14
  550. relationalai/early_access/metamodel/rewrite/__init__.py +0 -7
  551. relationalai/early_access/metamodel/typer/__init__.py +0 -3
  552. relationalai/early_access/metamodel/typer/typer/__init__.py +0 -12
  553. relationalai/early_access/metamodel/types/__init__.py +0 -15
  554. relationalai/early_access/metamodel/util/__init__.py +0 -15
  555. relationalai/early_access/metamodel/visitor/__init__.py +0 -12
  556. relationalai/early_access/rel/__init__.py +0 -12
  557. relationalai/early_access/rel/executor/__init__.py +0 -12
  558. relationalai/early_access/rel/rel_utils/__init__.py +0 -12
  559. relationalai/early_access/rel/rewrite/__init__.py +0 -7
  560. relationalai/early_access/solvers/__init__.py +0 -19
  561. relationalai/early_access/sql/__init__.py +0 -11
  562. relationalai/early_access/sql/executor/__init__.py +0 -3
  563. relationalai/early_access/sql/rewrite/__init__.py +0 -3
  564. relationalai/early_access/tests/logging/__init__.py +0 -12
  565. relationalai/early_access/tests/test_snapshot_base/__init__.py +0 -12
  566. relationalai/early_access/tests/utils/__init__.py +0 -12
  567. relationalai/environments/__init__.py +0 -35
  568. relationalai/environments/base.py +0 -381
  569. relationalai/environments/colab.py +0 -14
  570. relationalai/environments/generic.py +0 -71
  571. relationalai/environments/ipython.py +0 -68
  572. relationalai/environments/jupyter.py +0 -9
  573. relationalai/environments/snowbook.py +0 -169
  574. relationalai/errors.py +0 -2496
  575. relationalai/experimental/SF.py +0 -38
  576. relationalai/experimental/inspect.py +0 -47
  577. relationalai/experimental/pathfinder/__init__.py +0 -158
  578. relationalai/experimental/pathfinder/api.py +0 -160
  579. relationalai/experimental/pathfinder/automaton.py +0 -584
  580. relationalai/experimental/pathfinder/bridge.py +0 -226
  581. relationalai/experimental/pathfinder/compiler.py +0 -416
  582. relationalai/experimental/pathfinder/datalog.py +0 -214
  583. relationalai/experimental/pathfinder/diagnostics.py +0 -56
  584. relationalai/experimental/pathfinder/filter.py +0 -236
  585. relationalai/experimental/pathfinder/glushkov.py +0 -439
  586. relationalai/experimental/pathfinder/options.py +0 -265
  587. relationalai/experimental/pathfinder/pathfinder-v0.7.0.rel +0 -1951
  588. relationalai/experimental/pathfinder/rpq.py +0 -344
  589. relationalai/experimental/pathfinder/transition.py +0 -200
  590. relationalai/experimental/pathfinder/utils.py +0 -26
  591. relationalai/experimental/paths/README.md +0 -107
  592. relationalai/experimental/paths/api.py +0 -143
  593. relationalai/experimental/paths/benchmarks/grid_graph.py +0 -37
  594. relationalai/experimental/paths/code_organization.md +0 -2
  595. relationalai/experimental/paths/examples/Movies.ipynb +0 -16328
  596. relationalai/experimental/paths/examples/basic_example.py +0 -40
  597. relationalai/experimental/paths/examples/minimal_engine_warmup.py +0 -3
  598. relationalai/experimental/paths/examples/movie_example.py +0 -77
  599. relationalai/experimental/paths/examples/movies_data/actedin.csv +0 -193
  600. relationalai/experimental/paths/examples/movies_data/directed.csv +0 -45
  601. relationalai/experimental/paths/examples/movies_data/follows.csv +0 -7
  602. relationalai/experimental/paths/examples/movies_data/movies.csv +0 -39
  603. relationalai/experimental/paths/examples/movies_data/person.csv +0 -134
  604. relationalai/experimental/paths/examples/movies_data/produced.csv +0 -16
  605. relationalai/experimental/paths/examples/movies_data/ratings.csv +0 -10
  606. relationalai/experimental/paths/examples/movies_data/wrote.csv +0 -11
  607. relationalai/experimental/paths/examples/paths_benchmark.py +0 -115
  608. relationalai/experimental/paths/examples/paths_example.py +0 -116
  609. relationalai/experimental/paths/examples/pattern_to_automaton.py +0 -28
  610. relationalai/experimental/paths/find_paths_via_automaton.py +0 -85
  611. relationalai/experimental/paths/graph.py +0 -185
  612. relationalai/experimental/paths/path_algorithms/find_paths.py +0 -280
  613. relationalai/experimental/paths/path_algorithms/one_sided_ball_repetition.py +0 -26
  614. relationalai/experimental/paths/path_algorithms/one_sided_ball_upto.py +0 -111
  615. relationalai/experimental/paths/path_algorithms/single.py +0 -59
  616. relationalai/experimental/paths/path_algorithms/two_sided_balls_repetition.py +0 -39
  617. relationalai/experimental/paths/path_algorithms/two_sided_balls_upto.py +0 -103
  618. relationalai/experimental/paths/path_algorithms/usp-old.py +0 -130
  619. relationalai/experimental/paths/path_algorithms/usp-tuple.py +0 -183
  620. relationalai/experimental/paths/path_algorithms/usp.py +0 -150
  621. relationalai/experimental/paths/product_graph.py +0 -93
  622. relationalai/experimental/paths/rpq/automaton.py +0 -584
  623. relationalai/experimental/paths/rpq/diagnostics.py +0 -56
  624. relationalai/experimental/paths/rpq/rpq.py +0 -378
  625. relationalai/experimental/paths/tests/tests_limit_sp_max_length.py +0 -90
  626. relationalai/experimental/paths/tests/tests_limit_sp_multiple.py +0 -119
  627. relationalai/experimental/paths/tests/tests_limit_sp_single.py +0 -104
  628. relationalai/experimental/paths/tests/tests_limit_walks_multiple.py +0 -113
  629. relationalai/experimental/paths/tests/tests_limit_walks_single.py +0 -149
  630. relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_multiple.py +0 -70
  631. relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_single.py +0 -64
  632. relationalai/experimental/paths/tests/tests_one_sided_ball_upto_multiple.py +0 -115
  633. relationalai/experimental/paths/tests/tests_one_sided_ball_upto_single.py +0 -75
  634. relationalai/experimental/paths/tests/tests_single_paths.py +0 -152
  635. relationalai/experimental/paths/tests/tests_single_walks.py +0 -208
  636. relationalai/experimental/paths/tests/tests_single_walks_undirected.py +0 -297
  637. relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_multiple.py +0 -107
  638. relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_single.py +0 -76
  639. relationalai/experimental/paths/tests/tests_two_sided_balls_upto_multiple.py +0 -76
  640. relationalai/experimental/paths/tests/tests_two_sided_balls_upto_single.py +0 -110
  641. relationalai/experimental/paths/tests/tests_usp_nsp_multiple.py +0 -229
  642. relationalai/experimental/paths/tests/tests_usp_nsp_single.py +0 -108
  643. relationalai/experimental/paths/tree_agg.py +0 -168
  644. relationalai/experimental/paths/utilities/iterators.py +0 -27
  645. relationalai/experimental/paths/utilities/prefix_sum.py +0 -91
  646. relationalai/experimental/solvers.py +0 -1095
  647. relationalai/loaders/csv.py +0 -195
  648. relationalai/loaders/loader.py +0 -177
  649. relationalai/loaders/types.py +0 -23
  650. relationalai/rel_emitter.py +0 -373
  651. relationalai/rel_utils.py +0 -185
  652. relationalai/semantics/designs/query_builder/identify_by.md +0 -106
  653. relationalai/semantics/devtools/benchmark_lqp.py +0 -535
  654. relationalai/semantics/devtools/compilation_manager.py +0 -294
  655. relationalai/semantics/devtools/extract_lqp.py +0 -110
  656. relationalai/semantics/internal/internal.py +0 -3785
  657. relationalai/semantics/internal/snowflake.py +0 -329
  658. relationalai/semantics/lqp/README.md +0 -34
  659. relationalai/semantics/lqp/algorithms.py +0 -173
  660. relationalai/semantics/lqp/builtins.py +0 -213
  661. relationalai/semantics/lqp/compiler.py +0 -22
  662. relationalai/semantics/lqp/constructors.py +0 -68
  663. relationalai/semantics/lqp/executor.py +0 -518
  664. relationalai/semantics/lqp/export_rewriter.py +0 -40
  665. relationalai/semantics/lqp/intrinsics.py +0 -24
  666. relationalai/semantics/lqp/ir.py +0 -150
  667. relationalai/semantics/lqp/model2lqp.py +0 -1056
  668. relationalai/semantics/lqp/passes.py +0 -38
  669. relationalai/semantics/lqp/primitives.py +0 -252
  670. relationalai/semantics/lqp/result_helpers.py +0 -266
  671. relationalai/semantics/lqp/rewrite/__init__.py +0 -32
  672. relationalai/semantics/lqp/rewrite/algorithm.py +0 -385
  673. relationalai/semantics/lqp/rewrite/annotate_constraints.py +0 -69
  674. relationalai/semantics/lqp/rewrite/cdc.py +0 -216
  675. relationalai/semantics/lqp/rewrite/constants_to_vars.py +0 -70
  676. relationalai/semantics/lqp/rewrite/deduplicate_vars.py +0 -104
  677. relationalai/semantics/lqp/rewrite/eliminate_data.py +0 -108
  678. relationalai/semantics/lqp/rewrite/extract_common.py +0 -340
  679. relationalai/semantics/lqp/rewrite/extract_keys.py +0 -577
  680. relationalai/semantics/lqp/rewrite/flatten_script.py +0 -301
  681. relationalai/semantics/lqp/rewrite/function_annotations.py +0 -114
  682. relationalai/semantics/lqp/rewrite/functional_dependencies.py +0 -348
  683. relationalai/semantics/lqp/rewrite/period_math.py +0 -77
  684. relationalai/semantics/lqp/rewrite/quantify_vars.py +0 -339
  685. relationalai/semantics/lqp/rewrite/splinter.py +0 -76
  686. relationalai/semantics/lqp/rewrite/unify_definitions.py +0 -323
  687. relationalai/semantics/lqp/types.py +0 -101
  688. relationalai/semantics/lqp/utils.py +0 -170
  689. relationalai/semantics/lqp/validators.py +0 -70
  690. relationalai/semantics/metamodel/compiler.py +0 -134
  691. relationalai/semantics/metamodel/dependency.py +0 -880
  692. relationalai/semantics/metamodel/executor.py +0 -78
  693. relationalai/semantics/metamodel/factory.py +0 -287
  694. relationalai/semantics/metamodel/helpers.py +0 -368
  695. relationalai/semantics/metamodel/ir.py +0 -924
  696. relationalai/semantics/metamodel/rewrite/__init__.py +0 -8
  697. relationalai/semantics/metamodel/rewrite/discharge_constraints.py +0 -39
  698. relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py +0 -220
  699. relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py +0 -78
  700. relationalai/semantics/metamodel/rewrite/flatten.py +0 -590
  701. relationalai/semantics/metamodel/rewrite/format_outputs.py +0 -256
  702. relationalai/semantics/metamodel/rewrite/handle_aggregations_and_ranks.py +0 -237
  703. relationalai/semantics/metamodel/typer/checker.py +0 -355
  704. relationalai/semantics/metamodel/typer/typer.py +0 -1396
  705. relationalai/semantics/metamodel/util.py +0 -506
  706. relationalai/semantics/metamodel/visitor.py +0 -945
  707. relationalai/semantics/reasoners/__init__.py +0 -10
  708. relationalai/semantics/reasoners/graph/README.md +0 -620
  709. relationalai/semantics/reasoners/graph/__init__.py +0 -37
  710. relationalai/semantics/reasoners/graph/core.py +0 -9019
  711. relationalai/semantics/reasoners/graph/design/beyond_demand_transform.md +0 -797
  712. relationalai/semantics/reasoners/graph/tests/README.md +0 -21
  713. relationalai/semantics/reasoners/optimization/__init__.py +0 -68
  714. relationalai/semantics/reasoners/optimization/common.py +0 -88
  715. relationalai/semantics/reasoners/optimization/solvers_dev.py +0 -568
  716. relationalai/semantics/reasoners/optimization/solvers_pb.py +0 -1407
  717. relationalai/semantics/rel/builtins.py +0 -40
  718. relationalai/semantics/rel/compiler.py +0 -994
  719. relationalai/semantics/rel/executor.py +0 -363
  720. relationalai/semantics/rel/rel.py +0 -482
  721. relationalai/semantics/rel/rel_utils.py +0 -276
  722. relationalai/semantics/snowflake/__init__.py +0 -3
  723. relationalai/semantics/sql/compiler.py +0 -2503
  724. relationalai/semantics/sql/executor/duck_db.py +0 -52
  725. relationalai/semantics/sql/executor/result_helpers.py +0 -64
  726. relationalai/semantics/sql/executor/snowflake.py +0 -149
  727. relationalai/semantics/sql/rewrite/denormalize.py +0 -222
  728. relationalai/semantics/sql/rewrite/double_negation.py +0 -49
  729. relationalai/semantics/sql/rewrite/recursive_union.py +0 -127
  730. relationalai/semantics/sql/rewrite/sort_output_query.py +0 -246
  731. relationalai/semantics/sql/sql.py +0 -504
  732. relationalai/semantics/std/pragmas.py +0 -11
  733. relationalai/semantics/std/std.py +0 -14
  734. relationalai/semantics/tests/lqp/algorithms.py +0 -345
  735. relationalai/semantics/tests/test_snapshot_abstract.py +0 -144
  736. relationalai/semantics/tests/test_snapshot_base.py +0 -9
  737. relationalai/semantics/tests/utils.py +0 -46
  738. relationalai/std/__init__.py +0 -70
  739. relationalai/tools/cli.py +0 -2089
  740. relationalai/tools/cli_controls.py +0 -1975
  741. relationalai/tools/cli_helpers.py +0 -802
  742. relationalai/tools/debugger_client.py +0 -109
  743. relationalai/tools/debugger_server.py +0 -302
  744. relationalai/tools/dev.py +0 -685
  745. relationalai/tools/notes +0 -7
  746. relationalai/tools/qb_debugger.py +0 -425
  747. relationalai/tools/txn_progress.py +0 -188
  748. relationalai/util/clean_up_databases.py +0 -95
  749. relationalai/util/list_databases.py +0 -9
  750. relationalai/util/otel_configuration.py +0 -26
  751. relationalai/util/otel_handler.py +0 -484
  752. relationalai/util/snowflake_handler.py +0 -88
  753. relationalai/util/span_format_test.py +0 -43
  754. relationalai/util/span_tracker.py +0 -207
  755. relationalai/util/spans_file_handler.py +0 -72
  756. relationalai/util/tracing_handler.py +0 -34
  757. relationalai-0.13.5.dist-info/METADATA +0 -74
  758. relationalai-0.13.5.dist-info/RECORD +0 -473
  759. relationalai-0.13.5.dist-info/WHEEL +0 -4
  760. relationalai-0.13.5.dist-info/entry_points.txt +0 -3
  761. relationalai-0.13.5.dist-info/licenses/LICENSE +0 -202
  762. relationalai_test_util/__init__.py +0 -4
  763. relationalai_test_util/fixtures.py +0 -233
  764. relationalai_test_util/snapshot.py +0 -252
  765. relationalai_test_util/traceback.py +0 -118
  766. /relationalai/{analysis → semantics/frontend}/__init__.py +0 -0
  767. /relationalai/{auth/__init__.py → semantics/metamodel/metamodel_compiler.py} +0 -0
  768. /relationalai/{early_access → shims}/__init__.py +0 -0
  769. {relationalai/early_access/dsl/adapters → v0/relationalai/analysis}/__init__.py +0 -0
  770. {relationalai → v0/relationalai}/analysis/mechanistic.py +0 -0
  771. {relationalai → v0/relationalai}/analysis/whynot.py +0 -0
  772. {relationalai/early_access/dsl/adapters/orm → v0/relationalai/auth}/__init__.py +0 -0
  773. {relationalai → v0/relationalai}/auth/jwt_generator.py +0 -0
  774. {relationalai → v0/relationalai}/auth/oauth_callback_server.py +0 -0
  775. {relationalai → v0/relationalai}/auth/token_handler.py +0 -0
  776. {relationalai → v0/relationalai}/auth/util.py +0 -0
  777. {relationalai/clients/resources/snowflake → v0/relationalai/clients}/cache_store.py +0 -0
  778. {relationalai → v0/relationalai}/compiler.py +0 -0
  779. {relationalai → v0/relationalai}/dependencies.py +0 -0
  780. {relationalai → v0/relationalai}/docutils.py +0 -0
  781. {relationalai/early_access/dsl/adapters/owl → v0/relationalai/early_access}/__init__.py +0 -0
  782. {relationalai → v0/relationalai}/early_access/dsl/__init__.py +0 -0
  783. {relationalai/early_access/dsl/bindings → v0/relationalai/early_access/dsl/adapters}/__init__.py +0 -0
  784. {relationalai/early_access/dsl/bindings/legacy → v0/relationalai/early_access/dsl/adapters/orm}/__init__.py +0 -0
  785. {relationalai → v0/relationalai}/early_access/dsl/adapters/orm/model.py +0 -0
  786. {relationalai/early_access/dsl/codegen → v0/relationalai/early_access/dsl/adapters/owl}/__init__.py +0 -0
  787. {relationalai → v0/relationalai}/early_access/dsl/adapters/owl/model.py +0 -0
  788. {relationalai/early_access/dsl/core/temporal → v0/relationalai/early_access/dsl/bindings}/__init__.py +0 -0
  789. {relationalai/early_access/dsl/ir → v0/relationalai/early_access/dsl/bindings/legacy}/__init__.py +0 -0
  790. {relationalai/early_access/dsl/ontologies → v0/relationalai/early_access/dsl/codegen}/__init__.py +0 -0
  791. {relationalai → v0/relationalai}/early_access/dsl/constants.py +0 -0
  792. {relationalai → v0/relationalai}/early_access/dsl/core/__init__.py +0 -0
  793. {relationalai → v0/relationalai}/early_access/dsl/core/constraints/__init__.py +0 -0
  794. {relationalai → v0/relationalai}/early_access/dsl/core/constraints/predicate/__init__.py +0 -0
  795. {relationalai → v0/relationalai}/early_access/dsl/core/stack.py +0 -0
  796. {relationalai/early_access/dsl/orm → v0/relationalai/early_access/dsl/core/temporal}/__init__.py +0 -0
  797. {relationalai → v0/relationalai}/early_access/dsl/core/utils.py +0 -0
  798. {relationalai/early_access/dsl/orm/measures → v0/relationalai/early_access/dsl/ir}/__init__.py +0 -0
  799. {relationalai/early_access/dsl/physical_metadata → v0/relationalai/early_access/dsl/ontologies}/__init__.py +0 -0
  800. {relationalai → v0/relationalai}/early_access/dsl/ontologies/raw_source.py +0 -0
  801. {relationalai/early_access/dsl/serialize → v0/relationalai/early_access/dsl/orm}/__init__.py +0 -0
  802. {relationalai/early_access/dsl/snow → v0/relationalai/early_access/dsl/orm/measures}/__init__.py +0 -0
  803. {relationalai → v0/relationalai}/early_access/dsl/orm/reasoner_errors.py +0 -0
  804. {relationalai/loaders → v0/relationalai/early_access/dsl/physical_metadata}/__init__.py +0 -0
  805. {relationalai/semantics/tests → v0/relationalai/early_access/dsl/serialize}/__init__.py +0 -0
  806. {relationalai → v0/relationalai}/early_access/dsl/serialize/binding_model.py +0 -0
  807. {relationalai → v0/relationalai}/early_access/dsl/serialize/model.py +0 -0
  808. {relationalai/semantics/tests/lqp → v0/relationalai/early_access/dsl/snow}/__init__.py +0 -0
  809. {relationalai → v0/relationalai}/early_access/tests/__init__.py +0 -0
  810. {relationalai → v0/relationalai}/environments/ci.py +0 -0
  811. {relationalai → v0/relationalai}/environments/hex.py +0 -0
  812. {relationalai → v0/relationalai}/environments/terminal.py +0 -0
  813. {relationalai → v0/relationalai}/experimental/__init__.py +0 -0
  814. {relationalai → v0/relationalai}/experimental/graphs.py +0 -0
  815. {relationalai → v0/relationalai}/experimental/paths/__init__.py +0 -0
  816. {relationalai → v0/relationalai}/experimental/paths/benchmarks/__init__.py +0 -0
  817. {relationalai → v0/relationalai}/experimental/paths/path_algorithms/__init__.py +0 -0
  818. {relationalai → v0/relationalai}/experimental/paths/rpq/__init__.py +0 -0
  819. {relationalai → v0/relationalai}/experimental/paths/rpq/filter.py +0 -0
  820. {relationalai → v0/relationalai}/experimental/paths/rpq/glushkov.py +0 -0
  821. {relationalai → v0/relationalai}/experimental/paths/rpq/transition.py +0 -0
  822. {relationalai → v0/relationalai}/experimental/paths/utilities/__init__.py +0 -0
  823. {relationalai → v0/relationalai}/experimental/paths/utilities/utilities.py +0 -0
  824. {relationalai/tools → v0/relationalai/loaders}/__init__.py +0 -0
  825. {relationalai → v0/relationalai}/metagen.py +0 -0
  826. {relationalai → v0/relationalai}/metamodel.py +0 -0
  827. {relationalai → v0/relationalai}/rel.py +0 -0
  828. {relationalai → v0/relationalai}/semantics/devtools/__init__.py +0 -0
  829. {relationalai → v0/relationalai}/semantics/internal/__init__.py +0 -0
  830. {relationalai → v0/relationalai}/semantics/internal/annotations.py +0 -0
  831. {relationalai → v0/relationalai}/semantics/lqp/__init__.py +0 -0
  832. {relationalai → v0/relationalai}/semantics/lqp/pragmas.py +0 -0
  833. {relationalai → v0/relationalai}/semantics/metamodel/dataflow.py +0 -0
  834. {relationalai → v0/relationalai}/semantics/metamodel/typer/__init__.py +0 -0
  835. {relationalai → v0/relationalai}/semantics/metamodel/types.py +0 -0
  836. {relationalai → v0/relationalai}/semantics/reasoners/experimental/__init__.py +0 -0
  837. {relationalai → v0/relationalai}/semantics/rel/__init__.py +0 -0
  838. {relationalai → v0/relationalai}/semantics/sql/__init__.py +0 -0
  839. {relationalai → v0/relationalai}/semantics/sql/executor/__init__.py +0 -0
  840. {relationalai → v0/relationalai}/semantics/sql/rewrite/__init__.py +0 -0
  841. {relationalai → v0/relationalai}/semantics/tests/logging.py +0 -0
  842. {relationalai → v0/relationalai}/std/aggregates.py +0 -0
  843. {relationalai → v0/relationalai}/std/dates.py +0 -0
  844. {relationalai → v0/relationalai}/std/graphs.py +0 -0
  845. {relationalai → v0/relationalai}/std/inspect.py +0 -0
  846. {relationalai → v0/relationalai}/std/math.py +0 -0
  847. {relationalai → v0/relationalai}/std/re.py +0 -0
  848. {relationalai → v0/relationalai}/std/strings.py +0 -0
  849. {relationalai → v0/relationalai}/tools/cleanup_snapshots.py +0 -0
  850. {relationalai → v0/relationalai}/tools/constants.py +0 -0
  851. {relationalai → v0/relationalai}/tools/query_utils.py +0 -0
  852. {relationalai → v0/relationalai}/tools/snapshot_viewer.py +0 -0
  853. {relationalai → v0/relationalai}/util/__init__.py +0 -0
  854. {relationalai → v0/relationalai}/util/constants.py +0 -0
  855. {relationalai → v0/relationalai}/util/graph.py +0 -0
  856. {relationalai → v0/relationalai}/util/timeout.py +0 -0
@@ -1,1975 +0,0 @@
1
- #pyright: reportPrivateImportUsage=false
2
- from __future__ import annotations
3
-
4
- # Standard library imports
5
- import contextvars
6
- import io
7
- import itertools
8
- import os
9
- import shutil
10
- import sys
11
- import threading
12
- import time
13
- import importlib
14
- from dataclasses import dataclass
15
- from pathlib import Path
16
- from typing import Any, Callable, Dict, List, Optional, Sequence, TextIO, cast
17
-
18
- # Third-party imports
19
- import rich
20
- from InquirerPy import inquirer, utils as inquirer_utils
21
- from InquirerPy.base.complex import FakeDocument
22
- from InquirerPy.base.control import Choice
23
- from prompt_toolkit.key_binding import KeyPressEvent
24
- from prompt_toolkit.validation import ValidationError
25
- from rich.color import Color
26
- from rich.console import Console, Group
27
- from wcwidth import wcwidth
28
-
29
- # Local imports
30
- from relationalai import debugging
31
- from relationalai.util.format import format_duration
32
- from ..environments import (
33
- HexEnvironment,
34
- JupyterEnvironment,
35
- NotebookRuntimeEnvironment,
36
- SnowbookEnvironment,
37
- runtime_env,
38
- )
39
-
40
- # ---------------------------------------------
41
- # Global controls for nesting TaskProgress
42
- # ---------------------------------------------
43
-
44
- # Type alias for any progress type that supports nesting
45
- _ProgressType = Any # Actually TaskProgress | NotebookTaskProgress, but defined before those classes
46
-
47
- _current_progress: contextvars.ContextVar[Optional[_ProgressType]] = contextvars.ContextVar(
48
- 'current_progress', default=None
49
- )
50
-
51
-
52
- def get_current_progress() -> Optional[_ProgressType]:
53
- """Get the currently active TaskProgress, if any."""
54
- return _current_progress.get()
55
-
56
-
57
- def _set_current_progress(progress: Optional[_ProgressType]) -> contextvars.Token:
58
- """Set the current TaskProgress and return a token for restoration."""
59
- return _current_progress.set(progress)
60
-
61
- #--------------------------------------------------
62
- # Constants
63
- #--------------------------------------------------
64
-
65
- # Display symbols
66
- ARROW = "➜"
67
- CHECK_MARK = "✓"
68
- SUCCESS_ICON = "✅"
69
- FAIL_ICON = "❌"
70
-
71
- # Spinner animation frames
72
- SPINNER_FRAMES = ["▰▱▱▱", "▰▰▱▱", "▰▰▰▱", "▰▰▰▰", "▱▰▰▰", "▱▱▰▰", "▱▱▱▰", "▱▱▱▱"]
73
-
74
- # Terminal display constants
75
- DEFAULT_TERMINAL_WIDTH = 80
76
- SEPARATOR_WIDTH = 40
77
-
78
- # Task progress constants
79
- INITIALIZATION_COMPLETED_TEXT = "Parallel init finished in"
80
- MIN_CATEGORY_DURATION_SECONDS = 0.25 # Only show categories with duration > 250ms
81
-
82
- # Task category constants
83
- TASK_CATEGORY_INDEXING = "indexing"
84
- TASK_CATEGORY_PROVISIONING = "provisioning"
85
- TASK_CATEGORY_CHANGE_TRACKING = "change_tracking"
86
- TASK_CATEGORY_CACHE = "cache"
87
- TASK_CATEGORY_RELATIONS = "relations"
88
- TASK_CATEGORY_STATUS = "status"
89
- TASK_CATEGORY_VALIDATION = "validation"
90
- TASK_CATEGORY_OTHER = "other"
91
-
92
- # Default summary categories
93
- DEFAULT_SUMMARY_CATEGORIES = {
94
- TASK_CATEGORY_INDEXING: "Indexing",
95
- TASK_CATEGORY_PROVISIONING: "Provisioning",
96
- TASK_CATEGORY_CHANGE_TRACKING: "Change tracking",
97
- TASK_CATEGORY_RELATIONS: "Relations",
98
- TASK_CATEGORY_STATUS: "Status",
99
- TASK_CATEGORY_VALIDATION: "Validation",
100
- TASK_CATEGORY_OTHER: "Other"
101
- }
102
-
103
- # Parallel task categories (for duration calculation)
104
- PARALLEL_TASK_CATEGORIES = {
105
- TASK_CATEGORY_INDEXING,
106
- TASK_CATEGORY_PROVISIONING,
107
- TASK_CATEGORY_VALIDATION,
108
- TASK_CATEGORY_CHANGE_TRACKING
109
- }
110
-
111
- # Prompt constants
112
- REFETCH = "[REFETCH LIST]"
113
- MANUAL_ENTRY = "[MANUAL ENTRY]"
114
-
115
- # Timing constants
116
- HIGHLIGHT_DURATION = 2.0
117
- COMPLETION_DISPLAY_DURATION = 8.0
118
- TIMER_CHECK_INTERVAL = 0.1
119
- SPINNER_UPDATE_INTERVAL = 0.15
120
- INITIAL_DISPLAY_DELAY = 0.25
121
- BRIEF_PAUSE = 0.1
122
- LIVE_REFRESH_RATE = 10
123
-
124
- #--------------------------------------------------
125
- # Style
126
- #--------------------------------------------------
127
-
128
- STYLE = inquirer_utils.get_style({
129
- "fuzzy_prompt": "#e5c07b"
130
- }, False)
131
-
132
- #--------------------------------------------------
133
- # Helpers
134
- #--------------------------------------------------
135
-
136
- def rich_str(string:str, style:str|None = None) -> str:
137
- output = io.StringIO()
138
- console = Console(file=output, force_terminal=True)
139
- console.print(string, style=style)
140
- return output.getvalue()
141
-
142
- def _load_ipython_display() -> tuple[Any, Callable[..., Any]]:
143
- """Load IPython display helpers, raising if unavailable."""
144
- try:
145
- module = importlib.import_module("IPython.display")
146
- except ImportError as exc: # pragma: no cover - only triggered without IPython
147
- raise RuntimeError(
148
- "NotebookTaskProgress requires IPython when running in a notebook environment."
149
- ) from exc
150
-
151
- html_factory = getattr(module, "HTML")
152
- display_fn = getattr(module, "display")
153
- return html_factory, cast(Callable[..., Any], display_fn)
154
-
155
- def nat_path(path: Path, base: Path):
156
- resolved_path = path.resolve()
157
- resolved_base = base.resolve()
158
- if resolved_base in resolved_path.parents or resolved_path == resolved_base:
159
- return resolved_path.relative_to(resolved_base)
160
- else:
161
- return resolved_path.absolute()
162
-
163
- def get_default(value:str|None, list_of_values:Sequence[str]):
164
- if value is None:
165
- return None
166
- list_of_values_lower = [v.lower() for v in list_of_values]
167
- value_lower = value.lower()
168
- if value_lower in list_of_values_lower:
169
- return value
170
-
171
- #--------------------------------------------------
172
- # Dividers
173
- #--------------------------------------------------
174
-
175
- def divider(console=None, flush=False):
176
- div = "\n[dim]---------------------------------------------------\n "
177
- if console is None:
178
- rich.print(div)
179
- else:
180
- console.print(div)
181
- if flush:
182
- sys.stdout.flush()
183
-
184
- def abort():
185
- rich.print()
186
- rich.print("[yellow]Aborted")
187
- divider()
188
- sys.exit(1)
189
-
190
- #--------------------------------------------------
191
- # Prompts
192
- #--------------------------------------------------
193
-
194
- default_bindings = cast(Any, {
195
- "interrupt": [
196
- {"key": "escape"},
197
- {"key": "c-c"},
198
- {"key": "c-d"}
199
- ],
200
- "skip": [
201
- {"key": "c-s"}
202
- ]
203
- })
204
-
205
- def prompt(message:str, value:str|None, newline=False, validator:Callable|None = None, invalid_message:str|None = None) -> str:
206
- if value:
207
- return value
208
- if invalid_message is None:
209
- invalid_message = "Invalid input"
210
- try:
211
- result:str = inquirer.text(
212
- message,
213
- validate=validator,
214
- invalid_message=invalid_message,
215
- keybindings=default_bindings,
216
- ).execute()
217
- except KeyboardInterrupt:
218
- abort()
219
- raise Exception("Unreachable")
220
- if newline:
221
- rich.print("")
222
- return result
223
-
224
- def select(message:str, choices:List[str|Choice], value:str|None, newline=False, **kwargs) -> str|Any:
225
- if value:
226
- return value
227
- try:
228
- result:str = inquirer.select(message, choices, keybindings=default_bindings, **kwargs).execute()
229
- except KeyboardInterrupt:
230
- abort()
231
- raise Exception("Unreachable")
232
- if newline:
233
- rich.print("")
234
- return result
235
-
236
- def _enumerate_static_choices(choices: inquirer_utils.InquirerPyChoice) -> inquirer_utils.InquirerPyChoice:
237
- return [{"name": f"{i+1} {choice}", "value": choice} for i, choice in enumerate(choices)]
238
-
239
- def _enumerate_choices(choices: inquirer_utils.InquirerPyListChoices) -> inquirer_utils.InquirerPyListChoices:
240
- if callable(choices):
241
- return lambda session: _enumerate_static_choices(choices(session))
242
- else:
243
- return _enumerate_static_choices(choices)
244
-
245
- def _fuzzy(message:str, choices:inquirer_utils.InquirerPyListChoices, default:str|None = None, multiselect=False, show_index=False, **kwargs) -> str|list[str]|None:
246
- if show_index:
247
- choices = _enumerate_choices(choices)
248
-
249
- try:
250
- kwargs["keybindings"] = default_bindings
251
- if multiselect:
252
- kwargs["keybindings"] = { # pylint: disable=assignment-from-no-return
253
- "toggle": [
254
- {"key": "tab"}, # toggle choices
255
- ],
256
- "toggle-down": [
257
- {"key": "tab", "filter":False},
258
- ],
259
- }.update(default_bindings)
260
- kwargs["multiselect"] = True
261
-
262
- # NOTE: Using the builtin `default` kwarg to do this also filters
263
- # results which is undesirable and confusing for pre-filled
264
- # fields, so we move the cursor ourselves using the internals :(
265
- prompt = inquirer.fuzzy(message, choices=choices, max_height=8, border=True, style=STYLE, **kwargs)
266
- prompt._content_control._get_choices(prompt._content_control.choices, default)
267
-
268
- return prompt.execute()
269
- except KeyboardInterrupt:
270
- return abort()
271
-
272
- def fuzzy(message:str, choices:inquirer_utils.InquirerPyListChoices, default:str|None = None, show_index=False, **kwargs) -> str:
273
- return cast(str, _fuzzy(message, choices, default=default, show_index=show_index, **kwargs))
274
-
275
- def fuzzy_multiselect(message:str, choices:inquirer_utils.InquirerPyListChoices, default:str|None = None, show_index=False, **kwargs) -> list[str]:
276
- return cast(list[str], _fuzzy(message, choices, default=default, show_index=show_index, multiselect=True, **kwargs))
277
-
278
- def fuzzy_with_refetch(prompt: str, type: str, fn: Callable, *args, **kwargs):
279
- exception = None
280
- auto_select = kwargs.get("auto_select", None)
281
- not_found_message = kwargs.get("not_found_message", None)
282
- manual_entry = kwargs.get("manual_entry", None)
283
- items = []
284
- with Spinner(f"Fetching {type}", f"Fetched {type}"):
285
- try:
286
- items = fn(*args)
287
- except Exception as e:
288
- exception = e
289
- if exception is not None:
290
- rich.print(f"\n[red]Error fetching {type}: {exception}\n")
291
- return exception
292
- if len(items) == 0:
293
- if not_found_message:
294
- rich.print(f"\n[yellow]{not_found_message}\n")
295
- else:
296
- rich.print(f"\n[yellow]No valid {type} found\n")
297
- return None
298
-
299
- if auto_select and len(items) == 1 and items[0].lower() == auto_select.lower():
300
- return auto_select
301
-
302
- if manual_entry:
303
- items.insert(0, MANUAL_ENTRY)
304
- items.insert(0, REFETCH)
305
-
306
- passed_default = kwargs.get("default", None)
307
- passed_mandatory = kwargs.get("mandatory", False)
308
-
309
- rich.print("")
310
- result = fuzzy(
311
- prompt,
312
- items,
313
- default=get_default(passed_default, items),
314
- mandatory=passed_mandatory
315
- )
316
- rich.print("")
317
-
318
- while result == REFETCH:
319
- result = fuzzy_with_refetch(prompt, type, fn, *args, **kwargs)
320
- return result
321
-
322
- def confirm(message:str, default:bool = False) -> bool:
323
- try:
324
- return inquirer.confirm(message, default=default, keybindings=default_bindings).execute()
325
- except KeyboardInterrupt:
326
- return abort()
327
-
328
- def text(message:str, default:str|None = None, validator:Callable|None = None, invalid_message:str|None = None, **kwargs) -> str:
329
- if not invalid_message:
330
- invalid_message = "Invalid input"
331
- try:
332
- return inquirer.text(
333
- message,
334
- default=default or "",
335
- keybindings=default_bindings,
336
- validate=validator,
337
- invalid_message=invalid_message,
338
- **kwargs
339
- ).execute()
340
- except KeyboardInterrupt:
341
- return abort()
342
-
343
- def password(message:str, default:str|None = None, validator:Callable|None = None, invalid_message:str|None = None) -> str:
344
- if invalid_message is None:
345
- invalid_message = "Invalid input"
346
- try:
347
- return inquirer.secret(
348
- message,
349
- default=default or "",
350
- keybindings=default_bindings,
351
- validate=validator,
352
- invalid_message=invalid_message
353
- ).execute()
354
- except KeyboardInterrupt:
355
- return abort()
356
-
357
- def number(message:str, default:float|int|None = None, validator:Callable|None = None, invalid_message:str|None = None, **kwargs) -> float|int:
358
- if not invalid_message:
359
- invalid_message = "Invalid input"
360
- try:
361
- return inquirer.number(
362
- message,
363
- default=default or 0,
364
- keybindings=default_bindings,
365
- validate=validator,
366
- invalid_message=invalid_message,
367
- **kwargs
368
- ).execute()
369
- except KeyboardInterrupt:
370
- return abort()
371
-
372
- def file(message: str, start_path:Path|None = None, allow_freeform=False, **kwargs) -> str|None:
373
- try:
374
- return FuzzyFile(message, start_path, allow_freeform=allow_freeform, max_height=8, border=True, style=STYLE, **kwargs).execute()
375
- except KeyboardInterrupt:
376
- return abort()
377
-
378
- class FuzzyFile(inquirer.fuzzy):
379
- def __init__(self, message: str, initial_path: Path|None = None, allow_freeform = False, *args, **kwargs):
380
- self.initial_path = initial_path or Path()
381
- self.current_path = Path(self.initial_path)
382
- self.allow_freeform = allow_freeform
383
-
384
- kwargs["keybindings"] = {
385
- **default_bindings,
386
- "answer": [
387
- {"key": os.sep},
388
- {"key": "enter"},
389
- {"key": "tab"},
390
- {"key": "right"}
391
- ],
392
- **kwargs.get("keybindings", {})
393
- }
394
-
395
- super().__init__(message, *args, **kwargs, choices=self._get_choices)
396
-
397
- def _get_prompt_message(self) -> List[tuple[str, str]]:
398
- pre_answer = ("class:instruction", f" {self.instruction} " if self.instruction else " ")
399
- result = str(nat_path(self.current_path, self.initial_path))
400
-
401
- if result:
402
- sep = " " if self._amark else ""
403
- return [
404
- ("class:answermark", self._amark),
405
- ("class:answered_question", f"{sep}{self._message} "),
406
- ("class:answer", f"{result}{os.sep if not self.status['answered'] else ''}"),
407
- ]
408
- else:
409
- sep = " " if self._qmark else ""
410
- return [
411
- ("class:answermark", self._amark),
412
- ("class:questionmark", self._qmark),
413
- ("class:question", f"{sep}{self._message}"),
414
- pre_answer
415
- ]
416
-
417
- def _handle_enter(self, event: KeyPressEvent) -> None:
418
- try:
419
- fake_document = FakeDocument(self.result_value)
420
- self._validator.validate(fake_document) # type: ignore
421
- cc = self.content_control
422
- if self._multiselect:
423
- self.status["answered"] = True
424
- if not self.selected_choices:
425
- self.status["result"] = [cc.selection["name"]]
426
- event.app.exit(result=[cc.selection["value"]])
427
- else:
428
- self.status["result"] = self.result_name
429
- event.app.exit(result=self.result_value)
430
- else:
431
- res_value = cc.selection["value"]
432
- self.current_path /= res_value
433
- if self.current_path.is_dir():
434
- self._update_choices()
435
- else:
436
- self.status["answered"] = True
437
- self.status["result"] = cc.selection["name"]
438
- event.app.exit(result=str(nat_path(self.current_path, self.initial_path)))
439
- except ValidationError as e:
440
- self._set_error(str(e))
441
- except IndexError:
442
- self.status["answered"] = True
443
- res = self._get_current_text() if self.allow_freeform else None
444
- if self._multiselect:
445
- res = [res] if res is not None else []
446
- self.status["result"] = res
447
- event.app.exit(result=res)
448
-
449
- def _get_choices(self, _ = None):
450
- choices = os.listdir(self.current_path)
451
- choices.append("..")
452
- return choices
453
-
454
- def _update_choices(self):
455
- raw_choices = self._get_choices()
456
- cc = self.content_control
457
- cc.selected_choice_index = 0
458
- cc._raw_choices = raw_choices
459
- cc.choices = cc._get_choices(raw_choices, None)
460
- cc._safety_check()
461
- cc._format_choices()
462
- self._buffer.reset()
463
-
464
- #--------------------------------------------------
465
- # Line Clearing Mixin
466
- #--------------------------------------------------
467
-
468
- class LineClearingMixin:
469
- """Mixin class that provides line clearing functionality for different environments."""
470
-
471
- def __init__(self, *args, **kwargs):
472
- super().__init__(*args, **kwargs)
473
- self.last_line_length = 0
474
- # Detect environment capabilities
475
- import sys
476
- self.is_tty = sys.stdout.isatty()
477
- self.is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
478
- self.is_jupyter = isinstance(runtime_env, JupyterEnvironment)
479
-
480
- def _get_terminal_width(self):
481
- """Get terminal width, with fallback to reasonable default."""
482
- try:
483
- return shutil.get_terminal_size().columns
484
- except (OSError, AttributeError):
485
- return 80 # Fallback width
486
-
487
- def _clear_line(self, new_text: str):
488
- """Clear the current line and write new text using the best available method."""
489
- import sys
490
-
491
- if self.is_tty and not self.is_snowflake_notebook and not self.is_jupyter:
492
- # Use proper ANSI clear line sequence for terminals
493
- sys.stdout.write(f"\r\033[K{new_text}")
494
- else:
495
- # For notebooks and environments without ANSI support, use smart padding
496
- terminal_width = self._get_terminal_width()
497
-
498
- # Truncate text if it exceeds terminal width to prevent wrapping
499
- if len(new_text) > terminal_width:
500
- new_text = new_text[:terminal_width - 3] + "..."
501
-
502
- # Calculate how much of the line we need to clear
503
- # Use the maximum of last line length or terminal width to ensure full clearing
504
- clear_width = max(self.last_line_length, terminal_width)
505
-
506
- # Clear with spaces and write new text
507
- sys.stdout.write(f"\r{' ' * clear_width}\r{new_text}")
508
-
509
- sys.stdout.flush()
510
- # Update the tracked line length
511
- self.last_line_length = len(new_text)
512
-
513
- def _write_line(self, text: str, newline: bool = False):
514
- """Write text to the current line, optionally adding a newline."""
515
- import sys
516
- if newline:
517
- sys.stdout.write(f"{text}\n")
518
- else:
519
- sys.stdout.write(text)
520
- sys.stdout.flush()
521
-
522
- def _clear_and_write(self, text: str, newline: bool = False):
523
- """Clear the current line and write new text, with optional newline."""
524
- self._clear_line(text)
525
- if newline:
526
- import sys
527
- sys.stdout.write("\n")
528
- sys.stdout.flush()
529
-
530
-
531
- #--------------------------------------------------
532
- # Spinner
533
- #--------------------------------------------------
534
-
535
- class Spinner(LineClearingMixin):
536
- """Shows a spinner control while a task is running.
537
- The finished_message will not be printed if there was an exception and the failed_message is provided.
538
- """
539
- busy = False
540
-
541
- def __init__(
542
- self,
543
- message="",
544
- finished_message: str = "",
545
- failed_message=None,
546
- delay=None,
547
- leading_newline=False,
548
- trailing_newline=False,
549
- ):
550
- self.message = message
551
- self.finished_message = finished_message
552
- self.failed_message = failed_message
553
- self.spinner_generator = itertools.cycle(SPINNER_FRAMES)
554
- self.is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
555
- self.is_hex = isinstance(runtime_env, HexEnvironment)
556
- self.is_jupyter = isinstance(runtime_env, JupyterEnvironment)
557
- self.in_notebook = isinstance(runtime_env, NotebookRuntimeEnvironment)
558
- self.is_tty = sys.stdout.isatty()
559
-
560
- self._set_delay(delay)
561
- self.leading_newline = leading_newline
562
- self.trailing_newline = trailing_newline
563
- self.last_message = ""
564
- self.display = None
565
- # Add lock to prevent race conditions between spinner thread and main thread
566
- self._update_lock = threading.Lock()
567
-
568
- def _set_delay(self, delay: float|int|None) -> None:
569
- """Set appropriate delay based on environment and user input."""
570
- # If delay value is provided, validate and use it
571
- if delay:
572
- if isinstance(delay, (int, float)) and delay > 0:
573
- self.delay = float(delay)
574
- return
575
- else:
576
- raise ValueError(f"Invalid delay value: {delay}")
577
- # Otherwise, set delay based on environment
578
- elif self.is_hex:
579
- self.delay = 0 # Hex tries to append a new block each frame
580
- elif self.is_snowflake_notebook:
581
- self.delay = 0.5 # SF notebooks get bogged down
582
- elif self.in_notebook or self.is_tty:
583
- # Fast refresh for other notebooks or terminals with good printing support
584
- self.delay = 0.1
585
- else:
586
- # Otherwise disable the spinner animation entirely
587
- # for non-interactive environments.
588
- self.delay = 0
589
-
590
- def get_message(self, starting=False):
591
- max_width = shutil.get_terminal_size().columns
592
- spinner = "⏳⏳⏳⏳" if not self.is_tty and starting else next(self.spinner_generator)
593
- full_message = f"{spinner} {self.message}"
594
- if len(full_message) > max_width:
595
- return full_message[:max_width - 3] + "..."
596
- else:
597
- return full_message
598
-
599
- def update(self, message:str|None=None, color:str|None=None, file:TextIO|None=None, starting=False):
600
- # Use lock to prevent race conditions between spinner thread and main thread
601
- with self._update_lock:
602
- if message is None:
603
- message = self.get_message(starting=starting)
604
- if self.is_jupyter:
605
- # @NOTE: IPython isn't available in CI. This won't ever get invoked w/out IPython available though.
606
- from IPython.display import HTML, display # pyright: ignore[reportMissingImports]
607
- color_string = ""
608
- if color:
609
- color_value = Color.parse(color)
610
- rgb_tuple = color_value.get_truecolor()
611
- rgb_hex = f"#{rgb_tuple[0]:02X}{rgb_tuple[1]:02X}{rgb_tuple[2]:02X}"
612
- color_string = f"color: {rgb_hex};" if color is not None else ""
613
- content = HTML(f"<span style='font-family: monospace;{color_string}'>{message}</span>")
614
- if self.display is not None:
615
- self.display.update(content)
616
- else:
617
- self.display = display(content, display_id=True)
618
- else:
619
- if self.can_use_terminal_colors() and color is not None:
620
- rich_message = f"[{color}]{message}"
621
- else:
622
- rich_message = message
623
- rich_string = rich_str(rich_message)
624
- def width(word):
625
- return sum(wcwidth(c) for c in word)
626
- diff = width(self.last_message) - width(rich_string)
627
- self.reset_cursor()
628
- # Use rich.print with lock protection
629
- output_file = file or sys.stdout
630
- rich.print(rich_message + (" " * diff), file=output_file, end="", flush=False)
631
- if output_file.isatty() or self.in_notebook:
632
- output_file.flush()
633
- self.last_message = rich_string
634
-
635
- def can_use_terminal_colors(self):
636
- return not self.is_snowflake_notebook
637
-
638
- def update_messages(self, updater: dict[str, str]):
639
- if "message" in updater:
640
- self.message = updater["message"]
641
- if "finished_message" in updater:
642
- self.finished_message = updater["finished_message"]
643
- if "failed_message" in updater:
644
- self.failed_message = updater["failed_message"]
645
- self.update()
646
-
647
- def spinner_task(self):
648
- while self.busy and self.delay:
649
- self.update(color="magenta")
650
- time.sleep(self.delay) #type: ignore[union-attr] | we only call spinner_task if delay is not None anyway
651
- self.reset_cursor()
652
-
653
- def reset_cursor(self):
654
- if self.is_tty:
655
- # Clear the entire line and move cursor to beginning
656
- sys.stdout.write("\r\033[K")
657
- elif not self.is_jupyter:
658
- sys.stdout.write("\r")
659
-
660
- def __enter__(self):
661
- if self.leading_newline:
662
- rich.print()
663
- self.update(color="magenta", starting=True)
664
- # return control to the event loop briefly so stdout can be sure to flush:
665
- if self.delay:
666
- time.sleep(INITIAL_DISPLAY_DELAY)
667
- self.reset_cursor()
668
- if not self.delay:
669
- return self
670
- self.busy = True
671
- threading.Thread(target=self.spinner_task, daemon=True).start()
672
- return self
673
-
674
- def __exit__(self, exception, value, _):
675
- self.busy = False
676
- if exception is not None:
677
- if self.failed_message is not None:
678
- self.update(f"{self.failed_message} {value}", color="yellow", file=sys.stderr)
679
- # Use rich.print with explicit newline to ensure proper formatting
680
- rich.print(file=sys.stderr)
681
- return True
682
- return False
683
- if self.delay: # will be None for non-interactive environments
684
- time.sleep(self.delay)
685
- self.reset_cursor()
686
- if self.finished_message != "":
687
- final_message = f"▰▰▰▰ {self.finished_message}"
688
- self.update(final_message, color="green")
689
- # Use rich.print with explicit newline to ensure proper formatting
690
- rich.print()
691
- elif self.finished_message == "":
692
- self.update("")
693
- self.reset_cursor()
694
- if self.trailing_newline:
695
- rich.print()
696
-
697
- class DebuggingSpan:
698
- span: debugging.Span
699
- def __init__(self, span_type: str):
700
- self.span_type = span_type
701
- self.span_attrs = {}
702
-
703
- def attrs(self, **kwargs):
704
- self.span_attrs = kwargs
705
- return self
706
-
707
- def __enter__(self):
708
- self.span = debugging.span_start(self.span_type, **self.span_attrs)
709
- return self
710
-
711
- def __exit__(self, exc_type, exc_val, exc_tb):
712
- debugging.span_end(self.span)
713
-
714
- class SpanSpinner(Spinner):
715
- span: debugging.Span
716
- def __init__(self, span_type: str, *spinner_args, **spinner_kwargs):
717
- super().__init__(*spinner_args, **spinner_kwargs)
718
- self.span_type = span_type
719
- self.span_attrs = {}
720
-
721
- def attrs(self, **kwargs):
722
- self.span_attrs = kwargs
723
- return self
724
-
725
- def __enter__(self):
726
- self.span = debugging.span_start(self.span_type, **self.span_attrs)
727
- return super().__enter__()
728
-
729
- def __exit__(self, exc_type, exc_val, exc_tb):
730
- super().__exit__(exc_type, exc_val, exc_tb)
731
- debugging.span_end(self.span)
732
-
733
-
734
- @dataclass
735
- class TaskInfo:
736
- """Represents a single task with its state and metadata."""
737
- description: str
738
- category: str = "other"
739
- completed: bool = False
740
- added_time: float = 0.0
741
- completed_time: float = 0.0
742
- hidden: bool = False
743
-
744
- def __post_init__(self):
745
- if self.added_time == 0.0:
746
- self.added_time = time.time()
747
-
748
- def get_duration(self) -> float:
749
- """Get the duration of this task in seconds."""
750
- if not self.completed or self.completed_time == 0.0:
751
- return 0.0
752
-
753
- return self.completed_time - self.added_time
754
-
755
-
756
- class _TimerManager:
757
- """Manages all delayed operations for TaskProgress."""
758
-
759
- def __init__(self, progress_instance):
760
- self._progress = progress_instance
761
- self._operations = {} # task_id -> (operation_type, scheduled_time)
762
- self._thread = None
763
- self._running = False
764
-
765
- def schedule_highlight_removal(self, task_id: str, delay: float | None = None):
766
- """Schedule removal of highlighting for a task."""
767
- if delay is None:
768
- delay = HIGHLIGHT_DURATION
769
- scheduled_time = time.time() + delay
770
- self._operations[task_id] = ("remove_highlighting", scheduled_time)
771
- self._start()
772
-
773
- def schedule_task_removal(self, task_id: str, delay: float | None = None):
774
- """Schedule removal of a completed task."""
775
- if delay is None:
776
- delay = COMPLETION_DISPLAY_DURATION
777
- scheduled_time = time.time() + delay
778
- self._operations[task_id] = ("delayed_removal", scheduled_time)
779
- self._start()
780
-
781
- def schedule_task_hiding(self, task_id: str, delay: float | None = None):
782
- """Schedule hiding of a completed task from display (but keep in data structure)."""
783
- if delay is None:
784
- delay = COMPLETION_DISPLAY_DURATION
785
- scheduled_time = time.time() + delay
786
- self._operations[task_id] = ("delayed_hiding", scheduled_time)
787
- self._start()
788
-
789
-
790
- def _start(self):
791
- """Start the timer thread if not already running."""
792
- if self._thread is None or not self._thread.is_alive():
793
- self._running = True
794
- self._thread = threading.Thread(target=self._worker, daemon=True)
795
- self._thread.start()
796
-
797
- def _worker(self):
798
- """Worker thread for handling delayed operations."""
799
- while self._running:
800
- current_time = time.time()
801
- completed_ops = []
802
-
803
- # Find completed operations
804
- for task_id, (op_type, scheduled_time) in self._operations.items():
805
- if current_time >= scheduled_time:
806
- completed_ops.append((task_id, op_type))
807
-
808
- # Process completed operations
809
- for task_id, op_type in completed_ops:
810
- self._process_operation(task_id, op_type)
811
- del self._operations[task_id]
812
-
813
- time.sleep(TIMER_CHECK_INTERVAL)
814
-
815
- def _process_operation(self, task_id: str, op_type: str):
816
- """Process a completed delayed operation."""
817
- progress = self._progress
818
- if op_type == "remove_highlighting":
819
- if hasattr(progress, "_highlighted_tasks") and task_id in progress._highlighted_tasks:
820
- del progress._highlighted_tasks[task_id]
821
- if hasattr(progress, "_after_task_update"):
822
- progress._after_task_update()
823
- elif op_type == "delayed_removal":
824
- if hasattr(progress, "_tasks") and task_id in progress._tasks:
825
- del progress._tasks[task_id]
826
- if hasattr(progress, "_after_task_update"):
827
- progress._after_task_update()
828
- elif op_type == "delayed_hiding":
829
- if hasattr(progress, "_tasks") and task_id in progress._tasks:
830
- progress._tasks[task_id].hidden = True
831
- if hasattr(progress, "_after_task_update"):
832
- progress._after_task_update()
833
-
834
- def stop(self):
835
- """Stop the timer manager."""
836
- self._running = False
837
- self._operations.clear()
838
-
839
-
840
- class _TaskStateMixin:
841
- """Shared task management helpers for notebook and terminal progress displays."""
842
-
843
- enable_highlighting: bool = True
844
-
845
- def _init_task_state(self, *, hide_on_completion: bool = False, show_duration_summary: bool = True) -> None:
846
- self.hide_on_completion = hide_on_completion
847
- self.show_duration_summary = show_duration_summary
848
- self._tasks: dict[str, TaskInfo] = {}
849
- self._next_task_id: int = 1
850
- self._highlighted_tasks: dict[str, float] = {}
851
- self._process_start_time: float | None = None
852
- self._process_end_time: float | None = None
853
- self.main_completed: bool = False
854
- self.main_failed: bool = False
855
- self._timer_manager = _TimerManager(self)
856
-
857
- def _after_task_update(self) -> None:
858
- """Hook for subclasses to react when task state changes."""
859
- # Implemented by subclasses when they need to update the display immediately.
860
- return None
861
-
862
- def _generate_task_id(self) -> str:
863
- """Generate a unique task ID."""
864
- task_id = f"task_{self._next_task_id}"
865
- self._next_task_id += 1
866
- return task_id
867
-
868
- def add_sub_task(self, description: str, task_id: str | None = None, category: str = "general") -> str:
869
- """Add a new sub-task and return its unique ID."""
870
- if task_id is None:
871
- task_id = self._generate_task_id()
872
-
873
- if task_id not in self._tasks:
874
- self._tasks[task_id] = TaskInfo(description=description, category=category)
875
- self._after_task_update()
876
-
877
- return task_id
878
-
879
- def update_sub_task(self, task_id: str, description: str) -> None:
880
- """Update an existing sub-task description."""
881
- if task_id in self._tasks:
882
- task_info = self._tasks[task_id]
883
- if self.enable_highlighting and task_info.description != description:
884
- self._highlighted_tasks[task_id] = time.time() + HIGHLIGHT_DURATION
885
- self._timer_manager.schedule_highlight_removal(task_id)
886
-
887
- task_info.description = description
888
- self._after_task_update()
889
-
890
- def complete_sub_task(self, task_id: str, record_time: bool = True) -> None:
891
- """Complete a sub-task by marking it as done."""
892
- if task_id in self._tasks:
893
- if task_id in self._highlighted_tasks:
894
- del self._highlighted_tasks[task_id]
895
-
896
- if not self._tasks[task_id].completed and record_time:
897
- self._tasks[task_id].completed_time = time.time()
898
- self._tasks[task_id].completed = True
899
-
900
- self._after_task_update()
901
- self._timer_manager.schedule_task_hiding(task_id)
902
-
903
- def remove_sub_task(self, task_id: str, animate: bool = True) -> None:
904
- """Remove a sub-task by ID with optional completion animation."""
905
- if task_id in self._tasks:
906
- if task_id in self._highlighted_tasks:
907
- del self._highlighted_tasks[task_id]
908
-
909
- if animate:
910
- self.complete_sub_task(task_id)
911
- else:
912
- del self._tasks[task_id]
913
- self._after_task_update()
914
-
915
- def update_sub_status(self, sub_status: str) -> None:
916
- """Legacy method for backward compatibility - creates/updates a default sub-task."""
917
- self.add_sub_task(sub_status, "default")
918
- self.update_sub_task("default", sub_status)
919
-
920
- def update_main_status_fn(self, fn: Callable[[], str]) -> None:
921
- """Update the main status line using a callable function."""
922
- self._description_fn = fn
923
-
924
- def update_main_status(self, message: str) -> None:
925
- """Update the main status line with custom information."""
926
- if self._description_fn is not None:
927
- self._description_fn = None
928
- if getattr(self, "description", "") != message:
929
- self.description = message
930
- self._after_task_update()
931
-
932
- def update_messages(self, updater: dict[str, str]) -> None:
933
- """Update both main message and sub-status if provided."""
934
- if "message" in updater:
935
- if self._description_fn is not None:
936
- self._description_fn = None
937
- self.description = updater["message"]
938
- self._after_task_update()
939
- if "sub_status" in updater:
940
- self.update_sub_status(updater["sub_status"])
941
- if "success_message" in updater:
942
- self.success_message = updater["success_message"]
943
- if "failure_message" in updater:
944
- self.failure_message = updater["failure_message"]
945
-
946
- def get_sub_task_count(self) -> int:
947
- """Get the current number of active sub-tasks."""
948
- return len(self._tasks)
949
-
950
- def list_sub_tasks(self) -> list[str]:
951
- """Get a list of all active sub-task IDs."""
952
- return list(self._tasks.keys())
953
-
954
- def get_task_status(self) -> str:
955
- """Get a human-readable status of current task count vs limit."""
956
- return f"› Active tasks: {len(self._tasks)}"
957
-
958
- def get_task_duration(self, task_id: str) -> float:
959
- """Get the duration of a specific task in seconds."""
960
- if task_id in self._tasks:
961
- return self._tasks[task_id].get_duration()
962
- return 0.0
963
-
964
- def get_completed_tasks(self) -> dict[str, TaskInfo]:
965
- """Get all completed tasks with their timing information."""
966
- return {task_id: task_info for task_id, task_info in self._tasks.items() if task_info.completed}
967
-
968
- def get_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
969
- """Get all tasks (completed or active) for a specific category."""
970
- return {
971
- task_id: task_info
972
- for task_id, task_info in self._tasks.items()
973
- if task_info.category == category
974
- }
975
-
976
- def get_completed_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
977
- """Get all completed tasks for a specific category."""
978
- return {
979
- task_id: task_info
980
- for task_id, task_info in self._tasks.items()
981
- if task_info.category == category and task_info.completed
982
- }
983
-
984
- def _clear_all_tasks(self) -> None:
985
- """Clear all tasks and related data."""
986
- self._tasks.clear()
987
- self._highlighted_tasks.clear()
988
-
989
- def set_process_start_time(self) -> None:
990
- """Set the overall process start time."""
991
- self._process_start_time = time.time()
992
-
993
- def set_process_end_time(self) -> None:
994
- """Set the overall process end time."""
995
- self._process_end_time = time.time()
996
-
997
- def get_total_duration(self) -> float:
998
- """Get the total duration from first task added to last task completed."""
999
- if not self._tasks:
1000
- return 0.0
1001
-
1002
- completed_tasks = self.get_completed_tasks()
1003
- if not completed_tasks:
1004
- return 0.0
1005
-
1006
- start_times = [task.added_time for task in self._tasks.values()]
1007
- completion_times = [task.completed_time for task in completed_tasks.values() if task.completed_time > 0]
1008
-
1009
- if not start_times or not completion_times:
1010
- return 0.0
1011
-
1012
- earliest_start = min(start_times)
1013
- latest_completion = max(completion_times)
1014
- return latest_completion - earliest_start
1015
-
1016
-
1017
- class TaskProgress(_TaskStateMixin):
1018
- """A progress component that uses Rich's Live system to provide proper two-line display.
1019
-
1020
- This class provides:
1021
- - Main progress line with spinner and description
1022
- - Sub-status lines with hierarchical arrow indicators (➜)
1023
- - Proper error handling with success/failure messages
1024
- - Task-based progress tracking with context managers
1025
- - Highlighting of subtask text changes in yellow for 2 seconds when text differs
1026
- - Consistent task ordering with active tasks displayed above completed ones
1027
- """
1028
-
1029
- def __init__(
1030
- self,
1031
- description: str | Callable[[], str] = "",
1032
- success_message: str = "",
1033
- failure_message: str = "",
1034
- leading_newline: bool = False,
1035
- trailing_newline: bool = False,
1036
- transient: bool = False,
1037
- hide_on_completion: bool = False,
1038
- show_duration_summary: bool = True,
1039
- ):
1040
- # Public configuration - description can be a string or callable
1041
- if callable(description):
1042
- self._description_fn: Callable[[], str] | None = description
1043
- self.description = description() # Initial value
1044
- else:
1045
- self._description_fn = None
1046
- self.description = description
1047
- self.success_message = success_message
1048
- self.failure_message = failure_message
1049
- self.leading_newline = leading_newline
1050
- self.trailing_newline = trailing_newline
1051
- self.transient = transient
1052
- self._init_task_state(
1053
- hide_on_completion=hide_on_completion,
1054
- show_duration_summary=show_duration_summary,
1055
- )
1056
- self.enable_highlighting = True
1057
-
1058
- # Detect CI environment to avoid cursor control issues
1059
- from ..environments import CIEnvironment
1060
- self.is_ci = isinstance(runtime_env, CIEnvironment)
1061
- self.is_jupyter = isinstance(runtime_env, JupyterEnvironment)
1062
-
1063
- # Core components
1064
- # In CI or Jupyter, avoid forcing terminal rendering to prevent duplicate outputs
1065
- force_terminal = not self.is_ci and not self.is_jupyter
1066
- force_jupyter = True if self.is_jupyter else None
1067
- self.console = Console(
1068
- force_terminal=force_terminal,
1069
- force_jupyter=force_jupyter,
1070
- )
1071
- self.live = None
1072
- self.main_completed = False
1073
- self.main_failed = False
1074
-
1075
- # Animation state
1076
- self.spinner_index = 0
1077
-
1078
- # Performance optimizations
1079
- self._render_cache = None
1080
- self._last_state_hash = None
1081
-
1082
- self._spinner_thread = None
1083
-
1084
- # Nesting support
1085
- self._context_token: Optional[contextvars.Token] = None
1086
- self._parent_progress: Optional[TaskProgress] = None
1087
-
1088
- def _generate_task_id(self) -> str:
1089
- """Generate a unique task ID."""
1090
- task_id = f"task_{self._next_task_id}"
1091
- self._next_task_id += 1
1092
- return task_id
1093
-
1094
- def _get_description(self) -> str:
1095
- """Get the current description, calling the function if one was provided."""
1096
- if self._description_fn is not None:
1097
- return self._description_fn()
1098
- return self.description
1099
-
1100
- def _compute_state_hash(self, description: str) -> int:
1101
- """Compute a simple hash of the current state for caching."""
1102
- # Use a simple hash based on key state variables
1103
- state_parts = [
1104
- str(self.main_completed),
1105
- str(self.main_failed),
1106
- description,
1107
- str(self.spinner_index),
1108
- str(len(self._tasks)),
1109
- str(len(self._highlighted_tasks)),
1110
- ]
1111
-
1112
- # Add task states (only essential info for performance)
1113
- for task_id, task_info in self._tasks.items():
1114
- state_parts.append(f"{task_id}:{task_info.completed}:{task_info.description}")
1115
- if task_id in self._highlighted_tasks:
1116
- state_parts.append(f"highlight:{task_id}")
1117
-
1118
- return hash(tuple(state_parts))
1119
-
1120
- def _render_display(self):
1121
- """Render the current display state with caching optimization."""
1122
- # Get current description (may be dynamic)
1123
- description = self._get_description()
1124
-
1125
- # Check if we need to re-render
1126
- current_hash = self._compute_state_hash(description)
1127
- if current_hash == self._last_state_hash and self._render_cache is not None:
1128
- return self._render_cache
1129
-
1130
- from rich.text import Text
1131
-
1132
- # Build main task line
1133
- if self.main_failed:
1134
- # Split the description to style only the "Failed:" part in red
1135
- if description.startswith("Failed:"):
1136
- failed_part = "Failed:"
1137
- rest_part = description[len("Failed:"):].lstrip()
1138
- main_line = (Text(f"{FAIL_ICON} ", style="red") +
1139
- Text(failed_part, style="red") +
1140
- Text(f" {rest_part}", style="default"))
1141
- else:
1142
- # Fallback if description doesn't start with "Failed:"
1143
- main_line = Text(f"{FAIL_ICON} ", style="red") + Text(description, style="red")
1144
- elif self.main_completed:
1145
- main_line = Text(f"{SUCCESS_ICON} ", style="green") + Text(description, style="green")
1146
- else:
1147
- spinner_text = SPINNER_FRAMES[self.spinner_index]
1148
- main_line = Text(f"{spinner_text} ", style="magenta") + Text(description, style="magenta")
1149
-
1150
- # Build subtask lines
1151
- subtask_lines = self._render_subtask_lines()
1152
-
1153
- # Combine all lines
1154
- all_lines = [main_line] + subtask_lines
1155
-
1156
- # Cache the result
1157
- self._render_cache = Group(*all_lines)
1158
- self._last_state_hash = current_hash
1159
-
1160
- return self._render_cache
1161
-
1162
- def _render_subtask_lines(self):
1163
- """Render all subtask lines efficiently."""
1164
- from rich.text import Text
1165
-
1166
- subtask_lines = []
1167
- current_time = time.time()
1168
-
1169
- # Separate incomplete and completed tasks
1170
- incomplete_tasks = []
1171
- completed_tasks = []
1172
-
1173
- for task_id, task_info in self._tasks.items():
1174
- # Skip hidden tasks
1175
- if task_info.hidden:
1176
- continue
1177
- if task_info.completed:
1178
- completed_tasks.append((task_id, task_info))
1179
- else:
1180
- incomplete_tasks.append((task_id, task_info))
1181
-
1182
- # Render incomplete tasks first
1183
- for task_id, task_info in incomplete_tasks:
1184
- is_highlighted = (task_id in self._highlighted_tasks and
1185
- current_time < self._highlighted_tasks[task_id])
1186
-
1187
- style = "yellow" if is_highlighted else "white"
1188
- line = Text(f" {ARROW} ", style=style) + Text(task_info.description, style=style)
1189
- subtask_lines.append(line)
1190
-
1191
- # Render completed tasks
1192
- for task_id, task_info in completed_tasks:
1193
- line = Text(f" {CHECK_MARK} ", style="green") + Text(task_info.description, style="green")
1194
- subtask_lines.append(line)
1195
-
1196
- return subtask_lines
1197
-
1198
- def _advance_spinner(self):
1199
- """Advance the spinner animation."""
1200
- self.spinner_index = (self.spinner_index + 1) % len(SPINNER_FRAMES)
1201
-
1202
- def _invalidate_cache(self):
1203
- """Invalidate the render cache to force re-rendering."""
1204
- self._last_state_hash = None
1205
- self._render_cache = None
1206
-
1207
- def _update_display(self):
1208
- """Update the display if live."""
1209
- if self.live:
1210
- self.live.update(self._render_display())
1211
-
1212
- def _after_task_update(self) -> None:
1213
- """Refresh the live display when task state changes."""
1214
- self._invalidate_cache()
1215
- self._update_display()
1216
-
1217
- def _clear_all_tasks(self) -> None:
1218
- """Clear tasks and refresh display."""
1219
- super()._clear_all_tasks()
1220
- self._after_task_update()
1221
-
1222
- def generate_summary(self, categories: dict[str, str] | None = None) -> str:
1223
- """Generate a summary of completed tasks by category."""
1224
- if categories is None:
1225
- categories = DEFAULT_SUMMARY_CATEGORIES
1226
-
1227
- category_durations: dict[str, float] = {}
1228
- for category_name in categories:
1229
- tasks = self.get_completed_tasks_by_category(category_name)
1230
- category_durations[category_name] = _calculate_category_duration(category_name, tasks)
1231
-
1232
- if not any(category_durations.values()):
1233
- return ""
1234
-
1235
- total_duration = self.get_total_duration()
1236
-
1237
- try:
1238
- from rich.console import Console
1239
- from rich.table import Table
1240
-
1241
- table = Table(show_header=False, box=None, padding=(0, 1))
1242
- table.add_column("Operation", style="white")
1243
- table.add_column("Duration", style="green", justify="right")
1244
-
1245
- if total_duration > 0:
1246
- table.add_row(
1247
- INITIALIZATION_COMPLETED_TEXT,
1248
- format_duration(total_duration)
1249
- )
1250
-
1251
- for category_name, display_name in categories.items():
1252
- duration = category_durations[category_name]
1253
- if duration > MIN_CATEGORY_DURATION_SECONDS:
1254
- table.add_row(
1255
- f" {ARROW} {display_name}",
1256
- format_duration(duration)
1257
- )
1258
-
1259
- table.add_row("", "")
1260
-
1261
- console = Console()
1262
- with console.capture() as capture:
1263
- console.print(table)
1264
- return capture.get()
1265
-
1266
- except ImportError:
1267
- lines: list[str] = []
1268
- if total_duration > 0:
1269
- lines.append(f"{INITIALIZATION_COMPLETED_TEXT} {format_duration(total_duration)}")
1270
-
1271
- for category_name, display_name in categories.items():
1272
- duration = category_durations[category_name]
1273
- if duration > MIN_CATEGORY_DURATION_SECONDS:
1274
- lines.append(f" {ARROW} {display_name} {format_duration(duration)}")
1275
-
1276
- if lines:
1277
- lines.append("")
1278
-
1279
- return "\n".join(lines)
1280
-
1281
- def __enter__(self):
1282
- # Handle nesting: pause any parent progress
1283
- self._parent_progress = get_current_progress()
1284
- if self._parent_progress is not None:
1285
- self._parent_progress._pause()
1286
-
1287
- # Set ourselves as the current progress
1288
- self._context_token = _set_current_progress(self)
1289
-
1290
- if self.leading_newline:
1291
- print()
1292
-
1293
- # Start the live display
1294
- from rich.live import Live
1295
- self.live = Live(self._render_display(), console=self.console, refresh_per_second=LIVE_REFRESH_RATE)
1296
- self.live.start()
1297
-
1298
- # Start spinner animation
1299
- self._start_spinner()
1300
-
1301
- return self
1302
-
1303
- def _start_spinner(self):
1304
- """Start the spinner animation thread."""
1305
- self._spinner_paused = False
1306
-
1307
- def spinner_animation():
1308
- while self.live and not self.main_completed and not self.main_failed:
1309
- time.sleep(SPINNER_UPDATE_INTERVAL)
1310
- if self.live and not self._spinner_paused:
1311
- self._advance_spinner()
1312
- self.live.update(self._render_display())
1313
-
1314
- self._spinner_thread = threading.Thread(target=spinner_animation, daemon=True)
1315
- self._spinner_thread.start()
1316
-
1317
- def _pause(self):
1318
- """Pause the live display to allow a child progress to render."""
1319
- self._spinner_paused = True
1320
- if self.live:
1321
- # Clear the live display content before stopping so child has a clean slate
1322
- from rich.text import Text
1323
- self.live.update(Text(""))
1324
- self.live.stop()
1325
- # live.stop() leaves the cursor after the (empty) rendered content.
1326
- # Move cursor up one line so the child renders in the same place.
1327
- if not self.is_ci and sys.stdout.isatty():
1328
- sys.stdout.write("\033[A\r\033[K")
1329
- sys.stdout.flush()
1330
-
1331
- def _resume(self):
1332
- """Resume the live display after a child progress finishes."""
1333
- self._spinner_paused = False
1334
- if self.live and not self.main_completed and not self.main_failed:
1335
- # Re-render current state and restart the live display
1336
- self.live.update(self._render_display())
1337
- self.live.start()
1338
-
1339
- def __exit__(self, exc_type, exc_val, exc_tb):
1340
- # Stop timer manager
1341
- self._timer_manager.stop()
1342
-
1343
- if exc_type is not None:
1344
- # Exception occurred - show failure message
1345
- self._handle_failure(exc_val)
1346
- result = False # Don't suppress the exception
1347
- else:
1348
- # Success - show completion
1349
- self._handle_success()
1350
- result = True
1351
-
1352
- # Restore the parent progress as current
1353
- if self._context_token is not None:
1354
- _current_progress.reset(self._context_token)
1355
- self._context_token = None
1356
-
1357
- # Resume the parent progress if there was one
1358
- if self._parent_progress is not None:
1359
- self._parent_progress._resume()
1360
- self._parent_progress = None
1361
-
1362
- return result
1363
-
1364
- def _handle_failure(self, exc_val):
1365
- """Handle failure case in context manager exit."""
1366
- # Clear all tasks and update main task to show failure state
1367
- self._clear_all_tasks()
1368
- self.main_failed = True
1369
-
1370
- # Update main task description to show failure message
1371
- if self.failure_message:
1372
- self.description = self.failure_message
1373
- else:
1374
- self.description = f"Failed: {exc_val}"
1375
-
1376
- # Update the display to show the failure state before stopping
1377
- if self.live:
1378
- self.live.update(self._render_display())
1379
- # Brief pause to show the failure state
1380
- time.sleep(BRIEF_PAUSE)
1381
-
1382
- if self.trailing_newline:
1383
- print()
1384
- self._cleanup()
1385
-
1386
- def _handle_success(self):
1387
- """Handle success case in context manager exit."""
1388
- self.main_completed = True
1389
-
1390
- # Generate summary before clearing tasks (so we have the timing data)
1391
- # Only generate if show_duration_summary flag is True
1392
- summary = self.generate_summary() if self.show_duration_summary else ""
1393
-
1394
- self._clear_all_tasks()
1395
-
1396
- # Update main task description to show success message
1397
- if self.success_message:
1398
- self.description = self.success_message
1399
-
1400
- # Show success message in Rich Live display
1401
- if self.live:
1402
- self.live.update(self._render_display())
1403
- # Stop the live display
1404
- self.live.stop()
1405
-
1406
- # Print summary if available
1407
- if summary:
1408
- print() # Blank line for separation
1409
- print(summary, end="") # summary already has trailing newline
1410
- print() # Add extra blank line after summary
1411
-
1412
- if self.trailing_newline:
1413
- print()
1414
- self._cleanup()
1415
-
1416
- def _cleanup(self):
1417
- """Clean up resources."""
1418
- if self.live:
1419
- # Stop the live display first
1420
- self.live.stop()
1421
- # Clear the current line using ANSI escape sequence (only in TTY, not in CI)
1422
- if not self.is_ci and sys.stdout.isatty():
1423
- print("\r\033[K", end="", flush=True)
1424
-
1425
- def _calculate_category_duration(category_name: str, tasks: Dict[str, TaskInfo]) -> float:
1426
- """Calculate duration for a category based on task type (parallel vs sequential)."""
1427
- if not tasks:
1428
- return 0.0
1429
-
1430
- if category_name in PARALLEL_TASK_CATEGORIES:
1431
- # For parallel tasks, use time span (max completion - min start)
1432
- category_start_times = [task_info.added_time for task_info in tasks.values()]
1433
- category_completion_times = [
1434
- task_info.completed_time for task_info in tasks.values()
1435
- if task_info.completed_time > 0
1436
- ]
1437
- if category_start_times and category_completion_times:
1438
- return max(category_completion_times) - min(category_start_times)
1439
- else:
1440
- return 0.0
1441
- else:
1442
- # For sequential tasks, sum individual durations
1443
- return sum(task_info.get_duration() for task_info in tasks.values())
1444
-
1445
-
1446
- def create_progress(description: str | Callable[[], str] = "", success_message: str = "", failure_message: str = "",
1447
- leading_newline: bool = False, trailing_newline: bool = False, show_duration_summary: bool = True):
1448
- """Factory function to create the appropriate progress component based on environment.
1449
-
1450
- Automatically detects if we're in a notebook environment (Snowflake, Jupyter, etc.)
1451
- and returns the appropriate progress class.
1452
- """
1453
- from ..environments import runtime_env, SnowbookEnvironment, NotebookRuntimeEnvironment
1454
-
1455
- if isinstance(runtime_env, (SnowbookEnvironment, NotebookRuntimeEnvironment)):
1456
- # Use NotebookTaskProgress for Snowflake and Jupyter notebooks
1457
- return NotebookTaskProgress(
1458
- description=description,
1459
- success_message=success_message,
1460
- failure_message=failure_message,
1461
- leading_newline=leading_newline,
1462
- trailing_newline=trailing_newline,
1463
- show_duration_summary=show_duration_summary
1464
- )
1465
- else:
1466
- # Use TaskProgress for other environments (terminal, CI, etc.)
1467
- return TaskProgress(
1468
- description=description,
1469
- success_message=success_message,
1470
- failure_message=failure_message,
1471
- leading_newline=leading_newline,
1472
- trailing_newline=trailing_newline,
1473
- show_duration_summary=show_duration_summary
1474
- )
1475
-
1476
-
1477
- class SubTaskContext:
1478
- """Context manager for individual subtasks within a TaskProgress."""
1479
-
1480
- def __init__(self, task_progress: TaskProgress, description: str, task_id: str | None = None):
1481
- self.task_progress = task_progress
1482
- self.description = description
1483
- self.task_id = task_id
1484
- self._task_id = None
1485
-
1486
- def __enter__(self):
1487
- # Add the subtask and get its ID
1488
- self._task_id = self.task_progress.add_sub_task(self.description, self.task_id)
1489
- return self._task_id
1490
-
1491
- def __exit__(self, exc_type, exc_val, exc_tb):
1492
- if self._task_id and exc_type is None:
1493
- # Success - complete the subtask automatically when context exits
1494
- self.task_progress.complete_sub_task(self._task_id)
1495
- # If there was an exception, leave the subtask as-is for debugging
1496
- return False # Don't suppress exceptions
1497
-
1498
-
1499
- class NotebookTaskProgress(_TaskStateMixin):
1500
- """A progress component specifically designed for notebook environments like Snowflake.
1501
-
1502
- This class copies the EXACT working Spinner code and adapts it for notebook use.
1503
- """
1504
-
1505
- def __init__(
1506
- self,
1507
- description: str | Callable[[], str] = "",
1508
- success_message: str = "",
1509
- failure_message: str = "",
1510
- leading_newline: bool = False,
1511
- trailing_newline: bool = False,
1512
- show_duration_summary: bool = True,
1513
- ):
1514
- # Description can be a string or callable
1515
- if callable(description):
1516
- self._description_fn: Callable[[], str] | None = description
1517
- self.description = description() # Initial value
1518
- else:
1519
- self._description_fn = None
1520
- self.description = description
1521
- self.success_message = success_message
1522
- self.failure_message = failure_message
1523
- self.leading_newline = leading_newline
1524
- self.trailing_newline = trailing_newline
1525
- self._init_task_state(show_duration_summary=show_duration_summary)
1526
- self.enable_highlighting = False
1527
-
1528
- self.spinner_generator = itertools.cycle(SPINNER_FRAMES)
1529
-
1530
- # Environment detection for notebook environments only
1531
- self.is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
1532
- self.is_hex = isinstance(runtime_env, HexEnvironment)
1533
- self.is_jupyter = isinstance(runtime_env, JupyterEnvironment)
1534
- self.in_notebook = isinstance(runtime_env, NotebookRuntimeEnvironment)
1535
-
1536
- self._set_delay(None)
1537
-
1538
- self.last_message = ""
1539
- self.display = None
1540
- self._update_lock = threading.Lock()
1541
-
1542
- # Add sub-task support for TaskProgress compatibility
1543
- self.spinner_thread = None
1544
- self._current_subtask = ""
1545
- self.busy = False # Initialize busy state
1546
-
1547
- # Nesting support
1548
- self._context_token: Optional[contextvars.Token] = None
1549
- self._parent_progress: Optional[TaskProgress] = None
1550
- self._spinner_paused = False
1551
-
1552
- def _get_description(self) -> str:
1553
- """Get the current description, calling the function if one was provided."""
1554
- if self._description_fn is not None:
1555
- return self._description_fn()
1556
- return self.description
1557
-
1558
- def _generate_task_id(self) -> str:
1559
- """Generate a unique task ID."""
1560
- task_id = f"task_{self._next_task_id}"
1561
- self._next_task_id += 1
1562
- return task_id
1563
-
1564
- def _set_delay(self, delay: float|int|None) -> None:
1565
- """Set appropriate delay for notebook environments."""
1566
- # If delay value is provided, validate and use it
1567
- if delay:
1568
- if isinstance(delay, (int, float)) and delay > 0:
1569
- self.delay = float(delay)
1570
- return
1571
- else:
1572
- raise ValueError(f"Invalid delay value: {delay}")
1573
- # Simple delay for notebooks - no complex environment detection needed
1574
- elif self.in_notebook or self.is_snowflake_notebook or self.is_jupyter or self.is_hex:
1575
- self.delay = 0.2 # Simple, consistent delay for all notebook environments
1576
- else:
1577
- # Disable animation for non-interactive environments
1578
- self.delay = 0
1579
-
1580
- def get_message(self, starting=False):
1581
- """Get the current message with spinner - notebook environments only."""
1582
- # For notebook environments, use a reasonable default width
1583
- max_width = DEFAULT_TERMINAL_WIDTH
1584
- try:
1585
- max_width = shutil.get_terminal_size().columns
1586
- except (OSError, AttributeError):
1587
- pass # Use default width if terminal size can't be determined
1588
-
1589
- spinner = "⏳⏳⏳⏳" if starting else next(self.spinner_generator)
1590
-
1591
- # Get current description (may be dynamic)
1592
- description = self._get_description()
1593
-
1594
- # If there's an active subtask, show ONLY the subtask
1595
- if hasattr(self, '_current_subtask') and self._current_subtask:
1596
- full_message = f"{spinner} {self._current_subtask}"
1597
- else:
1598
- # Otherwise show the main task with subtask count if any
1599
- if len(self._tasks) > 0:
1600
- full_message = f"{spinner} {description} ({len(self._tasks)} active)"
1601
- else:
1602
- full_message = f"{spinner} {description}"
1603
-
1604
- if len(full_message) > max_width:
1605
- return full_message[:max_width - 3] + "..."
1606
- else:
1607
- return full_message
1608
-
1609
- def update(self, message:str|None=None, file:TextIO|None=None, starting=False):
1610
- """Update the display - notebook environments only."""
1611
- # Use lock to prevent race conditions between spinner thread and main thread
1612
- with self._update_lock:
1613
- if self.is_jupyter:
1614
- _, display_fn = _load_ipython_display()
1615
-
1616
- if message is None:
1617
- lines = self._build_jupyter_lines(starting=starting)
1618
- elif message == "":
1619
- lines = []
1620
- else:
1621
- lines = [message]
1622
-
1623
- rendered = "\n".join(lines)
1624
- content = {"text/plain": rendered}
1625
- if self.display is not None:
1626
- self.display.update(content, raw=True)
1627
- else:
1628
- self.display = display_fn(content, display_id=True, raw=True)
1629
- return
1630
-
1631
- if message is None:
1632
- message = self.get_message(starting=starting)
1633
-
1634
- rich_string = rich_str(message)
1635
-
1636
- def width(word):
1637
- return sum(wcwidth(c) for c in word)
1638
-
1639
- diff = width(self.last_message) - width(rich_string)
1640
-
1641
- sys.stdout.write("\r") # Move to beginning
1642
- sys.stdout.write(" " * DEFAULT_TERMINAL_WIDTH) # Clear with spaces
1643
- sys.stdout.write("\r") # Move back to beginning
1644
-
1645
- sys.stdout.write(message + (" " * diff)) # Write text directly
1646
- if self.in_notebook:
1647
- sys.stdout.flush() # Force output
1648
- self.last_message = rich_string
1649
-
1650
- def _build_jupyter_lines(self, starting: bool) -> list[str]:
1651
- """Compose the main status and subtasks for Jupyter display."""
1652
- description = self._get_description()
1653
- if self.busy or starting:
1654
- spinner = SPINNER_FRAMES[0] if starting else next(self.spinner_generator)
1655
- main_line = f"{spinner} {description}"
1656
- else:
1657
- main_text = self.success_message or description
1658
- main_line = f"{SUCCESS_ICON} {main_text}"
1659
-
1660
- visible_tasks = self._collect_visible_tasks()
1661
- lines = [main_line]
1662
-
1663
- for marker, tasks in (
1664
- (ARROW, visible_tasks["incomplete"]),
1665
- (CHECK_MARK, visible_tasks["completed"]),
1666
- ):
1667
- for task_info in tasks:
1668
- lines.append(f" {marker} {task_info.description}")
1669
-
1670
- return lines
1671
-
1672
- def _collect_visible_tasks(self) -> dict[str, list["TaskInfo"]]:
1673
- """Separate visible tasks into incomplete and completed lists."""
1674
- incomplete: list["TaskInfo"] = []
1675
- completed: list["TaskInfo"] = []
1676
-
1677
- for task_info in self._tasks.values():
1678
- if task_info.hidden:
1679
- continue
1680
- if task_info.completed:
1681
- completed.append(task_info)
1682
- else:
1683
- incomplete.append(task_info)
1684
-
1685
- return {"incomplete": incomplete, "completed": completed}
1686
-
1687
- def reset_cursor(self):
1688
- """Reset cursor to beginning of line - notebook environments only."""
1689
- # For notebook environments, use simple carriage return
1690
- if not self.is_jupyter:
1691
- sys.stdout.write("\r")
1692
-
1693
- def spinner_task(self):
1694
- """Spinner animation task."""
1695
- while self.busy and self.delay:
1696
- if not self._spinner_paused:
1697
- self.update()
1698
- time.sleep(self.delay) #type: ignore[union-attr] | we only call spinner_task if delay is not None anyway
1699
- self.reset_cursor()
1700
-
1701
- def _pause(self):
1702
- """Pause the display to allow a child progress to render."""
1703
- self._spinner_paused = True
1704
- # Clear the current line so child can render cleanly
1705
- if not self.is_jupyter:
1706
- sys.stdout.write("\r" + " " * DEFAULT_TERMINAL_WIDTH + "\r")
1707
- sys.stdout.flush()
1708
-
1709
- def _resume(self):
1710
- """Resume the display after a child progress finishes."""
1711
- self._spinner_paused = False
1712
- # Force an immediate update to restore display
1713
- if self.busy:
1714
- self.update()
1715
-
1716
- def _update_subtask_display(self, subtask_text: str):
1717
- """Update sub-task display - shows ONLY the subtask text."""
1718
- # Store the current display state
1719
- if not hasattr(self, '_current_display'):
1720
- self._current_display = ""
1721
-
1722
- # Only update if the display has changed
1723
- if self._current_display != subtask_text:
1724
- # Store the subtask text for the spinner to use
1725
- self._current_subtask = subtask_text
1726
- self._current_display = subtask_text
1727
- # The spinner will now show the subtask instead of main task
1728
-
1729
- def add_sub_task(self, description: str, task_id: str | None = None, category: str = "general") -> str:
1730
- task_id = super().add_sub_task(description, task_id, category)
1731
- # Update spinner display with the active subtask
1732
- if task_id in self._tasks:
1733
- self._update_subtask_display(self._tasks[task_id].description)
1734
- return task_id
1735
-
1736
- def update_sub_task(self, task_id: str, description: str) -> None:
1737
- super().update_sub_task(task_id, description)
1738
- if task_id in self._tasks:
1739
- self._update_subtask_display(description)
1740
-
1741
- def complete_sub_task(self, task_id: str, record_time: bool = True) -> None:
1742
- super().complete_sub_task(task_id, record_time=record_time)
1743
- # Clear the subtask display when completed
1744
- self._current_subtask = ""
1745
- self._current_display = ""
1746
-
1747
- def remove_sub_task(self, task_id: str, animate: bool = True) -> None:
1748
- """Remove a sub-task by ID."""
1749
- task_description: str | None = None
1750
- if task_id in self._tasks:
1751
- task_description = self._tasks[task_id].description
1752
-
1753
- # Notebook display should drop the subtask immediately to clear the UI.
1754
- super().remove_sub_task(task_id, animate=False)
1755
-
1756
- if task_description and getattr(self, "_current_subtask", "") == task_description:
1757
- self._current_subtask = ""
1758
- self._current_display = ""
1759
- # The spinner will now show the main task again
1760
-
1761
- def update_sub_status(self, sub_status: str):
1762
- """Legacy method for backward compatibility - creates/updates a default sub-task."""
1763
- super().update_sub_status(sub_status)
1764
-
1765
- def update_main_status(self, message: str):
1766
- """Update the main status line with real-time updates."""
1767
- super().update_main_status(message)
1768
- self._current_subtask = ""
1769
- self._current_display = ""
1770
- # The spinner will now show the updated main task
1771
-
1772
- def update_messages(self, updater: dict[str, str]):
1773
- """Update both main message and sub-status if provided."""
1774
- super().update_messages(updater)
1775
-
1776
- def _clear_all_tasks(self) -> None:
1777
- super()._clear_all_tasks()
1778
- self._current_subtask = ""
1779
- self._current_display = ""
1780
-
1781
- def __enter__(self):
1782
- # Handle nesting: pause any parent progress
1783
- self._parent_progress = get_current_progress()
1784
- if self._parent_progress is not None:
1785
- self._parent_progress._pause()
1786
-
1787
- # Set ourselves as the current progress
1788
- self._context_token = _set_current_progress(self)
1789
-
1790
- # Skip leading newline for Jupyter - it interferes with IPython display
1791
- if self.leading_newline and not self.is_jupyter:
1792
- rich.print()
1793
- self.update(starting=True)
1794
- # return control to the event loop briefly so stdout can be sure to flush:
1795
- if self.delay:
1796
- time.sleep(INITIAL_DISPLAY_DELAY)
1797
- self.reset_cursor()
1798
- if not self.delay:
1799
- return self
1800
- self.busy = True
1801
- threading.Thread(target=self.spinner_task, daemon=True).start()
1802
- return self
1803
-
1804
- def __exit__(self, exc_type, exc_val, exc_tb):
1805
- self.busy = False
1806
- if exc_type is not None:
1807
- if self.failure_message is not None:
1808
- self.update(f"{self.failure_message} {exc_val}", file=sys.stderr)
1809
- # For non-Jupyter, add newline to ensure proper formatting
1810
- # For Jupyter, IPython display handles formatting
1811
- if not self.is_jupyter:
1812
- rich.print(file=sys.stderr)
1813
- result = True
1814
- else:
1815
- result = False
1816
- else:
1817
- if self.delay: # will be None for non-interactive environments
1818
- time.sleep(self.delay)
1819
-
1820
- # Generate summary BEFORE clearing the spinner line (so we have timing data)
1821
- # Only generate if show_duration_summary flag is True
1822
- summary = self.generate_summary() if self.show_duration_summary else ""
1823
-
1824
- # Clear the spinner line completely
1825
- self._clear_spinner_line()
1826
-
1827
- final_message: str | None = None
1828
- if self.success_message:
1829
- final_message = f"{SUCCESS_ICON} {self.success_message}"
1830
- elif summary:
1831
- final_message = f"{SUCCESS_ICON} Done"
1832
-
1833
- if final_message:
1834
- if self.is_jupyter:
1835
- if self.display is not None:
1836
- self.display.update({"text/plain": final_message}, raw=True)
1837
- else:
1838
- _, display_fn = _load_ipython_display()
1839
- self.display = display_fn({"text/plain": final_message}, display_id=True, raw=True)
1840
- else:
1841
- print(final_message)
1842
- elif self.success_message == "":
1843
- # When there's no success message, clear the display for notebooks
1844
- # The summary will be printed below if available
1845
- if self.is_jupyter:
1846
- self.update("")
1847
- # For non-Jupyter notebooks, _clear_spinner_line() already handled it
1848
-
1849
- # Print summary if there are completed tasks
1850
- if summary:
1851
- if self.is_jupyter:
1852
- # Use IPython display to avoid blank stdout lines in notebooks
1853
- _, display_fn = _load_ipython_display()
1854
- display_fn({"text/plain": summary.strip()}, raw=True)
1855
- else:
1856
- print()
1857
- print(summary.strip()) # Summary includes visual separator line
1858
-
1859
- # Skip trailing newline for Jupyter - it interferes with IPython display
1860
- if self.trailing_newline and not self.is_jupyter:
1861
- rich.print()
1862
- result = True
1863
-
1864
- # Restore the parent progress as current
1865
- if self._context_token is not None:
1866
- _current_progress.reset(self._context_token)
1867
- self._context_token = None
1868
-
1869
- # Resume the parent progress if there was one
1870
- if self._parent_progress is not None:
1871
- self._parent_progress._resume()
1872
- self._parent_progress = None
1873
-
1874
- return result
1875
-
1876
- def _clear_spinner_line(self):
1877
- """Clear the current spinner line completely."""
1878
- # Skip clearing for Jupyter notebooks - IPython display handles it
1879
- if self.is_jupyter:
1880
- return
1881
-
1882
- # Write enough spaces to clear any content, then move to start of line
1883
- terminal_width = DEFAULT_TERMINAL_WIDTH
1884
- try:
1885
- terminal_width = shutil.get_terminal_size().columns
1886
- except (OSError, AttributeError):
1887
- pass
1888
-
1889
- # Clear with spaces, carriage return, and newline to ensure we're on a fresh line
1890
- sys.stdout.write("\r" + " " * terminal_width + "\r\n")
1891
- sys.stdout.flush()
1892
-
1893
- def set_process_start_time(self) -> None:
1894
- """Set the overall process start time."""
1895
- self._process_start_time = time.time()
1896
-
1897
- def set_process_end_time(self) -> None:
1898
- """Set the overall process end time."""
1899
- self._process_end_time = time.time()
1900
-
1901
- def get_total_duration(self) -> float:
1902
- """Get the total duration from first task added to last task completed."""
1903
- if not self._tasks:
1904
- return 0.0
1905
-
1906
- completed_tasks = self.get_completed_tasks()
1907
- if not completed_tasks:
1908
- return 0.0
1909
-
1910
- # Find earliest start time and latest completion time
1911
- start_times = [task.added_time for task in self._tasks.values()]
1912
- completion_times = [task.completed_time for task in completed_tasks.values() if task.completed_time > 0]
1913
-
1914
- if not start_times or not completion_times:
1915
- return 0.0
1916
-
1917
- earliest_start = min(start_times)
1918
- latest_completion = max(completion_times)
1919
-
1920
- return latest_completion - earliest_start
1921
-
1922
- def generate_summary(self, categories: dict[str, str] | None = None) -> str:
1923
- """Generate a summary of completed tasks by category."""
1924
- if categories is None:
1925
- categories = DEFAULT_SUMMARY_CATEGORIES
1926
-
1927
- # Get completed tasks by category and calculate durations
1928
- category_durations = {}
1929
- for category_name in categories:
1930
- tasks = self.get_completed_tasks_by_category(category_name)
1931
- category_durations[category_name] = _calculate_category_duration(category_name, tasks)
1932
-
1933
- # If there's nothing meaningful to show, return empty string
1934
- if not any(category_durations.values()):
1935
- return ""
1936
-
1937
- # Generate summary lines with proper alignment
1938
- summary_lines = []
1939
- label_width = 30 # Width for category labels
1940
- time_width = 10 # Width for time column (right-aligned)
1941
-
1942
- # Add total time FIRST (at the top) - align with arrow lines
1943
- total_duration = self.get_total_duration()
1944
- if total_duration > 0:
1945
- formatted_total = format_duration(total_duration)
1946
- # Use the same format as arrow lines but with a different prefix
1947
- # This ensures perfect alignment with the time column
1948
- summary_lines.append(f" {INITIALIZATION_COMPLETED_TEXT:<{label_width-1}} {formatted_total:>{time_width}}")
1949
-
1950
- # Add category breakdown
1951
- category_lines = []
1952
- for category_name, display_name in categories.items():
1953
- duration = category_durations[category_name]
1954
- if duration > MIN_CATEGORY_DURATION_SECONDS: # Only show significant durations
1955
- formatted_duration = format_duration(duration)
1956
- # Use arrow for visual consistency with right-aligned time
1957
- category_lines.append(f" {ARROW} {display_name:<{label_width-4}} {formatted_duration:>{time_width}}")
1958
-
1959
- # Only add category lines if there are any
1960
- if category_lines:
1961
- summary_lines.extend(category_lines)
1962
-
1963
- # Add a visual separator line for Snowflake notebook environment
1964
- summary_lines.append("─" * SEPARATOR_WIDTH)
1965
-
1966
- return "\n".join(summary_lines) + "\n"
1967
-
1968
- def get_completed_tasks(self) -> dict[str, TaskInfo]:
1969
- """Get all completed tasks with their timing information."""
1970
- return {task_id: task_info for task_id, task_info in self._tasks.items() if task_info.completed}
1971
-
1972
- def get_completed_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
1973
- """Get all completed tasks for a specific category."""
1974
- return {task_id: task_info for task_id, task_info in self._tasks.items()
1975
- if task_info.category == category and task_info.completed}