sequence.py 45 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653
  1. """Collection of string classes and utilities."""
  2. from __future__ import annotations
  3. import dataclasses
  4. import inspect
  5. import json
  6. import re
  7. import sys
  8. import typing
  9. from typing import (
  10. TYPE_CHECKING,
  11. Any,
  12. ClassVar,
  13. Dict,
  14. List,
  15. Literal,
  16. NoReturn,
  17. Set,
  18. Tuple,
  19. Type,
  20. TypeVar,
  21. Union,
  22. overload,
  23. )
  24. from reflex import constants
  25. from reflex.constants.base import REFLEX_VAR_OPENING_TAG
  26. from reflex.utils.exceptions import VarTypeError
  27. from reflex.utils.types import GenericType, get_origin
  28. from .base import (
  29. CachedVarOperation,
  30. CustomVarOperationReturn,
  31. LiteralNoneVar,
  32. LiteralVar,
  33. ToOperation,
  34. Var,
  35. VarData,
  36. _global_vars,
  37. cached_property_no_lock,
  38. figure_out_type,
  39. get_unique_variable_name,
  40. unionize,
  41. var_operation,
  42. var_operation_return,
  43. )
  44. from .number import (
  45. BooleanVar,
  46. LiteralNumberVar,
  47. NumberVar,
  48. raise_unsupported_operand_types,
  49. )
  50. if TYPE_CHECKING:
  51. from .object import ObjectVar
  52. class StringVar(Var[str]):
  53. """Base class for immutable string vars."""
  54. @overload
  55. def __add__(self, other: StringVar | str) -> ConcatVarOperation: ...
  56. @overload
  57. def __add__(self, other: NoReturn) -> NoReturn: ...
  58. def __add__(self, other: Any) -> ConcatVarOperation:
  59. """Concatenate two strings.
  60. Args:
  61. other: The other string.
  62. Returns:
  63. The string concatenation operation.
  64. """
  65. if not isinstance(other, (StringVar, str)):
  66. raise_unsupported_operand_types("+", (type(self), type(other)))
  67. return ConcatVarOperation.create(self, other)
  68. @overload
  69. def __radd__(self, other: StringVar | str) -> ConcatVarOperation: ...
  70. @overload
  71. def __radd__(self, other: NoReturn) -> NoReturn: ...
  72. def __radd__(self, other: Any) -> ConcatVarOperation:
  73. """Concatenate two strings.
  74. Args:
  75. other: The other string.
  76. Returns:
  77. The string concatenation operation.
  78. """
  79. if not isinstance(other, (StringVar, str)):
  80. raise_unsupported_operand_types("+", (type(other), type(self)))
  81. return ConcatVarOperation.create(other, self)
  82. @overload
  83. def __mul__(self, other: NumberVar | int) -> StringVar: ...
  84. @overload
  85. def __mul__(self, other: NoReturn) -> NoReturn: ...
  86. def __mul__(self, other: Any) -> StringVar:
  87. """Multiply the sequence by a number or an integer.
  88. Args:
  89. other: The number or integer to multiply the sequence by.
  90. Returns:
  91. StringVar: The resulting sequence after multiplication.
  92. """
  93. if not isinstance(other, (NumberVar, int)):
  94. raise_unsupported_operand_types("*", (type(self), type(other)))
  95. return (self.split() * other).join()
  96. @overload
  97. def __rmul__(self, other: NumberVar | int) -> StringVar: ...
  98. @overload
  99. def __rmul__(self, other: NoReturn) -> NoReturn: ...
  100. def __rmul__(self, other: Any) -> StringVar:
  101. """Multiply the sequence by a number or an integer.
  102. Args:
  103. other: The number or integer to multiply the sequence by.
  104. Returns:
  105. StringVar: The resulting sequence after multiplication.
  106. """
  107. if not isinstance(other, (NumberVar, int)):
  108. raise_unsupported_operand_types("*", (type(other), type(self)))
  109. return (self.split() * other).join()
  110. @overload
  111. def __getitem__(self, i: slice) -> StringVar: ...
  112. @overload
  113. def __getitem__(self, i: int | NumberVar) -> StringVar: ...
  114. def __getitem__(self, i: Any) -> StringVar:
  115. """Get a slice of the string.
  116. Args:
  117. i: The slice.
  118. Returns:
  119. The string slice operation.
  120. """
  121. if isinstance(i, slice):
  122. return self.split()[i].join()
  123. if not isinstance(i, (int, NumberVar)) or (
  124. isinstance(i, NumberVar) and i._is_strict_float()
  125. ):
  126. raise_unsupported_operand_types("[]", (type(self), type(i)))
  127. return string_item_operation(self, i)
  128. def length(self) -> NumberVar:
  129. """Get the length of the string.
  130. Returns:
  131. The string length operation.
  132. """
  133. return self.split().length()
  134. def lower(self) -> StringVar:
  135. """Convert the string to lowercase.
  136. Returns:
  137. The string lower operation.
  138. """
  139. return string_lower_operation(self)
  140. def upper(self) -> StringVar:
  141. """Convert the string to uppercase.
  142. Returns:
  143. The string upper operation.
  144. """
  145. return string_upper_operation(self)
  146. def strip(self) -> StringVar:
  147. """Strip the string.
  148. Returns:
  149. The string strip operation.
  150. """
  151. return string_strip_operation(self)
  152. def reversed(self) -> StringVar:
  153. """Reverse the string.
  154. Returns:
  155. The string reverse operation.
  156. """
  157. return self.split().reverse().join()
  158. @overload
  159. def contains(
  160. self, other: StringVar | str, field: StringVar | str | None = None
  161. ) -> BooleanVar: ...
  162. @overload
  163. def contains(
  164. self, other: NoReturn, field: StringVar | str | None = None
  165. ) -> NoReturn: ...
  166. def contains(self, other: Any, field: Any = None) -> BooleanVar:
  167. """Check if the string contains another string.
  168. Args:
  169. other: The other string.
  170. field: The field to check.
  171. Returns:
  172. The string contains operation.
  173. """
  174. if not isinstance(other, (StringVar, str)):
  175. raise_unsupported_operand_types("contains", (type(self), type(other)))
  176. if field is not None:
  177. if not isinstance(field, (StringVar, str)):
  178. raise_unsupported_operand_types("contains", (type(self), type(field)))
  179. return string_contains_field_operation(self, other, field)
  180. return string_contains_operation(self, other)
  181. @overload
  182. def split(self, separator: StringVar | str = "") -> ArrayVar[List[str]]: ...
  183. @overload
  184. def split(self, separator: NoReturn) -> NoReturn: ...
  185. def split(self, separator: Any = "") -> ArrayVar[List[str]]:
  186. """Split the string.
  187. Args:
  188. separator: The separator.
  189. Returns:
  190. The string split operation.
  191. """
  192. if not isinstance(separator, (StringVar, str)):
  193. raise_unsupported_operand_types("split", (type(self), type(separator)))
  194. return string_split_operation(self, separator)
  195. @overload
  196. def startswith(self, prefix: StringVar | str) -> BooleanVar: ...
  197. @overload
  198. def startswith(self, prefix: NoReturn) -> NoReturn: ...
  199. def startswith(self, prefix: Any) -> BooleanVar:
  200. """Check if the string starts with a prefix.
  201. Args:
  202. prefix: The prefix.
  203. Returns:
  204. The string starts with operation.
  205. """
  206. if not isinstance(prefix, (StringVar, str)):
  207. raise_unsupported_operand_types("startswith", (type(self), type(prefix)))
  208. return string_starts_with_operation(self, prefix)
  209. @overload
  210. def __lt__(self, other: StringVar | str) -> BooleanVar: ...
  211. @overload
  212. def __lt__(self, other: NoReturn) -> NoReturn: ...
  213. def __lt__(self, other: Any):
  214. """Check if the string is less than another string.
  215. Args:
  216. other: The other string.
  217. Returns:
  218. The string less than operation.
  219. """
  220. if not isinstance(other, (StringVar, str)):
  221. raise_unsupported_operand_types("<", (type(self), type(other)))
  222. return string_lt_operation(self, other)
  223. @overload
  224. def __gt__(self, other: StringVar | str) -> BooleanVar: ...
  225. @overload
  226. def __gt__(self, other: NoReturn) -> NoReturn: ...
  227. def __gt__(self, other: Any):
  228. """Check if the string is greater than another string.
  229. Args:
  230. other: The other string.
  231. Returns:
  232. The string greater than operation.
  233. """
  234. if not isinstance(other, (StringVar, str)):
  235. raise_unsupported_operand_types(">", (type(self), type(other)))
  236. return string_gt_operation(self, other)
  237. @overload
  238. def __le__(self, other: StringVar | str) -> BooleanVar: ...
  239. @overload
  240. def __le__(self, other: NoReturn) -> NoReturn: ...
  241. def __le__(self, other: Any):
  242. """Check if the string is less than or equal to another string.
  243. Args:
  244. other: The other string.
  245. Returns:
  246. The string less than or equal operation.
  247. """
  248. if not isinstance(other, (StringVar, str)):
  249. raise_unsupported_operand_types("<=", (type(self), type(other)))
  250. return string_le_operation(self, other)
  251. @overload
  252. def __ge__(self, other: StringVar | str) -> BooleanVar: ...
  253. @overload
  254. def __ge__(self, other: NoReturn) -> NoReturn: ...
  255. def __ge__(self, other: Any):
  256. """Check if the string is greater than or equal to another string.
  257. Args:
  258. other: The other string.
  259. Returns:
  260. The string greater than or equal operation.
  261. """
  262. if not isinstance(other, (StringVar, str)):
  263. raise_unsupported_operand_types(">=", (type(self), type(other)))
  264. return string_ge_operation(self, other)
  265. @var_operation
  266. def string_lt_operation(lhs: StringVar | str, rhs: StringVar | str):
  267. """Check if a string is less than another string.
  268. Args:
  269. lhs: The left-hand side string.
  270. rhs: The right-hand side string.
  271. Returns:
  272. The string less than operation.
  273. """
  274. return var_operation_return(js_expression=f"{lhs} < {rhs}", var_type=bool)
  275. @var_operation
  276. def string_gt_operation(lhs: StringVar | str, rhs: StringVar | str):
  277. """Check if a string is greater than another string.
  278. Args:
  279. lhs: The left-hand side string.
  280. rhs: The right-hand side string.
  281. Returns:
  282. The string greater than operation.
  283. """
  284. return var_operation_return(js_expression=f"{lhs} > {rhs}", var_type=bool)
  285. @var_operation
  286. def string_le_operation(lhs: StringVar | str, rhs: StringVar | str):
  287. """Check if a string is less than or equal to another string.
  288. Args:
  289. lhs: The left-hand side string.
  290. rhs: The right-hand side string.
  291. Returns:
  292. The string less than or equal operation.
  293. """
  294. return var_operation_return(js_expression=f"{lhs} <= {rhs}", var_type=bool)
  295. @var_operation
  296. def string_ge_operation(lhs: StringVar | str, rhs: StringVar | str):
  297. """Check if a string is greater than or equal to another string.
  298. Args:
  299. lhs: The left-hand side string.
  300. rhs: The right-hand side string.
  301. Returns:
  302. The string greater than or equal operation.
  303. """
  304. return var_operation_return(js_expression=f"{lhs} >= {rhs}", var_type=bool)
  305. @var_operation
  306. def string_lower_operation(string: StringVar):
  307. """Convert a string to lowercase.
  308. Args:
  309. string: The string to convert.
  310. Returns:
  311. The lowercase string.
  312. """
  313. return var_operation_return(js_expression=f"{string}.toLowerCase()", var_type=str)
  314. @var_operation
  315. def string_upper_operation(string: StringVar):
  316. """Convert a string to uppercase.
  317. Args:
  318. string: The string to convert.
  319. Returns:
  320. The uppercase string.
  321. """
  322. return var_operation_return(js_expression=f"{string}.toUpperCase()", var_type=str)
  323. @var_operation
  324. def string_strip_operation(string: StringVar):
  325. """Strip a string.
  326. Args:
  327. string: The string to strip.
  328. Returns:
  329. The stripped string.
  330. """
  331. return var_operation_return(js_expression=f"{string}.trim()", var_type=str)
  332. @var_operation
  333. def string_contains_field_operation(
  334. haystack: StringVar, needle: StringVar | str, field: StringVar | str
  335. ):
  336. """Check if a string contains another string.
  337. Args:
  338. haystack: The haystack.
  339. needle: The needle.
  340. field: The field to check.
  341. Returns:
  342. The string contains operation.
  343. """
  344. return var_operation_return(
  345. js_expression=f"{haystack}.some(obj => obj[{field}] === {needle})",
  346. var_type=bool,
  347. )
  348. @var_operation
  349. def string_contains_operation(haystack: StringVar, needle: StringVar | str):
  350. """Check if a string contains another string.
  351. Args:
  352. haystack: The haystack.
  353. needle: The needle.
  354. Returns:
  355. The string contains operation.
  356. """
  357. return var_operation_return(
  358. js_expression=f"{haystack}.includes({needle})", var_type=bool
  359. )
  360. @var_operation
  361. def string_starts_with_operation(full_string: StringVar, prefix: StringVar | str):
  362. """Check if a string starts with a prefix.
  363. Args:
  364. full_string: The full string.
  365. prefix: The prefix.
  366. Returns:
  367. Whether the string starts with the prefix.
  368. """
  369. return var_operation_return(
  370. js_expression=f"{full_string}.startsWith({prefix})", var_type=bool
  371. )
  372. @var_operation
  373. def string_item_operation(string: StringVar, index: NumberVar | int):
  374. """Get an item from a string.
  375. Args:
  376. string: The string.
  377. index: The index of the item.
  378. Returns:
  379. The item from the string.
  380. """
  381. return var_operation_return(js_expression=f"{string}.at({index})", var_type=str)
  382. @var_operation
  383. def array_join_operation(array: ArrayVar, sep: StringVar | str = ""):
  384. """Join the elements of an array.
  385. Args:
  386. array: The array.
  387. sep: The separator.
  388. Returns:
  389. The joined elements.
  390. """
  391. return var_operation_return(js_expression=f"{array}.join({sep})", var_type=str)
  392. # Compile regex for finding reflex var tags.
  393. _decode_var_pattern_re = (
  394. rf"{constants.REFLEX_VAR_OPENING_TAG}(.*?){constants.REFLEX_VAR_CLOSING_TAG}"
  395. )
  396. _decode_var_pattern = re.compile(_decode_var_pattern_re, flags=re.DOTALL)
  397. @dataclasses.dataclass(
  398. eq=False,
  399. frozen=True,
  400. **{"slots": True} if sys.version_info >= (3, 10) else {},
  401. )
  402. class LiteralStringVar(LiteralVar, StringVar):
  403. """Base class for immutable literal string vars."""
  404. _var_value: str = dataclasses.field(default="")
  405. @classmethod
  406. def create(
  407. cls,
  408. value: str,
  409. _var_type: GenericType | None = str,
  410. _var_data: VarData | None = None,
  411. ) -> StringVar:
  412. """Create a var from a string value.
  413. Args:
  414. value: The value to create the var from.
  415. _var_type: The type of the var.
  416. _var_data: Additional hooks and imports associated with the Var.
  417. Returns:
  418. The var.
  419. """
  420. if REFLEX_VAR_OPENING_TAG in value:
  421. strings_and_vals: list[Var | str] = []
  422. offset = 0
  423. # Find all tags
  424. while m := _decode_var_pattern.search(value):
  425. start, end = m.span()
  426. strings_and_vals.append(value[:start])
  427. serialized_data = m.group(1)
  428. if serialized_data.isnumeric() or (
  429. serialized_data[0] == "-" and serialized_data[1:].isnumeric()
  430. ):
  431. # This is a global immutable var.
  432. var = _global_vars[int(serialized_data)]
  433. strings_and_vals.append(var)
  434. value = value[(end + len(var._js_expr)) :]
  435. offset += end - start
  436. strings_and_vals.append(value)
  437. filtered_strings_and_vals = [
  438. s for s in strings_and_vals if isinstance(s, Var) or s
  439. ]
  440. if len(filtered_strings_and_vals) == 1:
  441. only_string = filtered_strings_and_vals[0]
  442. if isinstance(only_string, str):
  443. return LiteralVar.create(only_string).to(StringVar, _var_type)
  444. else:
  445. return only_string.to(StringVar, only_string._var_type)
  446. if len(
  447. literal_strings := [
  448. s
  449. for s in filtered_strings_and_vals
  450. if isinstance(s, (str, LiteralStringVar))
  451. ]
  452. ) == len(filtered_strings_and_vals):
  453. return LiteralStringVar.create(
  454. "".join(
  455. s._var_value if isinstance(s, LiteralStringVar) else s
  456. for s in literal_strings
  457. ),
  458. _var_type=_var_type,
  459. _var_data=VarData.merge(
  460. _var_data,
  461. *(
  462. s._get_all_var_data()
  463. for s in filtered_strings_and_vals
  464. if isinstance(s, Var)
  465. ),
  466. ),
  467. )
  468. concat_result = ConcatVarOperation.create(
  469. *filtered_strings_and_vals,
  470. _var_data=_var_data,
  471. )
  472. return (
  473. concat_result
  474. if _var_type is str
  475. else concat_result.to(StringVar, _var_type)
  476. )
  477. return LiteralStringVar(
  478. _js_expr=json.dumps(value),
  479. _var_type=_var_type,
  480. _var_data=_var_data,
  481. _var_value=value,
  482. )
  483. def __hash__(self) -> int:
  484. """Get the hash of the var.
  485. Returns:
  486. The hash of the var.
  487. """
  488. return hash((self.__class__.__name__, self._var_value))
  489. def json(self) -> str:
  490. """Get the JSON representation of the var.
  491. Returns:
  492. The JSON representation of the var.
  493. """
  494. return json.dumps(self._var_value)
  495. @dataclasses.dataclass(
  496. eq=False,
  497. frozen=True,
  498. **{"slots": True} if sys.version_info >= (3, 10) else {},
  499. )
  500. class ConcatVarOperation(CachedVarOperation, StringVar):
  501. """Representing a concatenation of literal string vars."""
  502. _var_value: Tuple[Var, ...] = dataclasses.field(default_factory=tuple)
  503. @cached_property_no_lock
  504. def _cached_var_name(self) -> str:
  505. """The name of the var.
  506. Returns:
  507. The name of the var.
  508. """
  509. list_of_strs: List[Union[str, Var]] = []
  510. last_string = ""
  511. for var in self._var_value:
  512. if isinstance(var, LiteralStringVar):
  513. last_string += var._var_value
  514. else:
  515. if last_string:
  516. list_of_strs.append(last_string)
  517. last_string = ""
  518. list_of_strs.append(var)
  519. if last_string:
  520. list_of_strs.append(last_string)
  521. list_of_strs_filtered = [
  522. str(LiteralVar.create(s)) for s in list_of_strs if isinstance(s, Var) or s
  523. ]
  524. if len(list_of_strs_filtered) == 1:
  525. return list_of_strs_filtered[0]
  526. return "(" + "+".join(list_of_strs_filtered) + ")"
  527. @cached_property_no_lock
  528. def _cached_get_all_var_data(self) -> VarData | None:
  529. """Get all the VarData asVarDatae Var.
  530. Returns:
  531. The VarData associated with the Var.
  532. """
  533. return VarData.merge(
  534. *[
  535. var._get_all_var_data()
  536. for var in self._var_value
  537. if isinstance(var, Var)
  538. ],
  539. self._var_data,
  540. )
  541. @classmethod
  542. def create(
  543. cls,
  544. *value: Var | str,
  545. _var_data: VarData | None = None,
  546. ) -> ConcatVarOperation:
  547. """Create a var from a string value.
  548. Args:
  549. value: The values to concatenate.
  550. _var_data: Additional hooks and imports associated with the Var.
  551. Returns:
  552. The var.
  553. """
  554. return cls(
  555. _js_expr="",
  556. _var_type=str,
  557. _var_data=_var_data,
  558. _var_value=tuple(map(LiteralVar.create, value)),
  559. )
  560. ARRAY_VAR_TYPE = TypeVar("ARRAY_VAR_TYPE", bound=Union[List, Tuple, Set])
  561. OTHER_TUPLE = TypeVar("OTHER_TUPLE")
  562. INNER_ARRAY_VAR = TypeVar("INNER_ARRAY_VAR")
  563. KEY_TYPE = TypeVar("KEY_TYPE")
  564. VALUE_TYPE = TypeVar("VALUE_TYPE")
  565. class ArrayVar(Var[ARRAY_VAR_TYPE]):
  566. """Base class for immutable array vars."""
  567. @overload
  568. def join(self, sep: StringVar | str = "") -> StringVar: ...
  569. @overload
  570. def join(self, sep: NoReturn) -> NoReturn: ...
  571. def join(self, sep: Any = "") -> StringVar:
  572. """Join the elements of the array.
  573. Args:
  574. sep: The separator between elements.
  575. Returns:
  576. The joined elements.
  577. """
  578. if not isinstance(sep, (StringVar, str)):
  579. raise_unsupported_operand_types("join", (type(self), type(sep)))
  580. if (
  581. isinstance(self, LiteralArrayVar)
  582. and (
  583. len(
  584. args := [
  585. x
  586. for x in self._var_value
  587. if isinstance(x, (LiteralStringVar, str))
  588. ]
  589. )
  590. == len(self._var_value)
  591. )
  592. and isinstance(sep, (LiteralStringVar, str))
  593. ):
  594. sep_str = sep._var_value if isinstance(sep, LiteralStringVar) else sep
  595. return LiteralStringVar.create(
  596. sep_str.join(
  597. i._var_value if isinstance(i, LiteralStringVar) else i for i in args
  598. )
  599. )
  600. return array_join_operation(self, sep)
  601. def reverse(self) -> ArrayVar[ARRAY_VAR_TYPE]:
  602. """Reverse the array.
  603. Returns:
  604. The reversed array.
  605. """
  606. return array_reverse_operation(self)
  607. @overload
  608. def __add__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> ArrayVar[ARRAY_VAR_TYPE]: ...
  609. @overload
  610. def __add__(self, other: NoReturn) -> NoReturn: ...
  611. def __add__(self, other: Any) -> ArrayVar[ARRAY_VAR_TYPE]:
  612. """Concatenate two arrays.
  613. Parameters:
  614. other: The other array to concatenate.
  615. Returns:
  616. ArrayConcatOperation: The concatenation of the two arrays.
  617. """
  618. if not isinstance(other, ArrayVar):
  619. raise_unsupported_operand_types("+", (type(self), type(other)))
  620. return array_concat_operation(self, other)
  621. @overload
  622. def __getitem__(self, i: slice) -> ArrayVar[ARRAY_VAR_TYPE]: ...
  623. @overload
  624. def __getitem__(
  625. self: (
  626. ArrayVar[Tuple[int, OTHER_TUPLE]]
  627. | ArrayVar[Tuple[float, OTHER_TUPLE]]
  628. | ArrayVar[Tuple[int | float, OTHER_TUPLE]]
  629. ),
  630. i: Literal[0, -2],
  631. ) -> NumberVar: ...
  632. @overload
  633. def __getitem__(
  634. self: (
  635. ArrayVar[Tuple[OTHER_TUPLE, int]]
  636. | ArrayVar[Tuple[OTHER_TUPLE, float]]
  637. | ArrayVar[Tuple[OTHER_TUPLE, int | float]]
  638. ),
  639. i: Literal[1, -1],
  640. ) -> NumberVar: ...
  641. @overload
  642. def __getitem__(
  643. self: ArrayVar[Tuple[str, OTHER_TUPLE]], i: Literal[0, -2]
  644. ) -> StringVar: ...
  645. @overload
  646. def __getitem__(
  647. self: ArrayVar[Tuple[OTHER_TUPLE, str]], i: Literal[1, -1]
  648. ) -> StringVar: ...
  649. @overload
  650. def __getitem__(
  651. self: ArrayVar[Tuple[bool, OTHER_TUPLE]], i: Literal[0, -2]
  652. ) -> BooleanVar: ...
  653. @overload
  654. def __getitem__(
  655. self: ArrayVar[Tuple[OTHER_TUPLE, bool]], i: Literal[1, -1]
  656. ) -> BooleanVar: ...
  657. @overload
  658. def __getitem__(
  659. self: (
  660. ARRAY_VAR_OF_LIST_ELEMENT[int]
  661. | ARRAY_VAR_OF_LIST_ELEMENT[float]
  662. | ARRAY_VAR_OF_LIST_ELEMENT[int | float]
  663. ),
  664. i: int | NumberVar,
  665. ) -> NumberVar: ...
  666. @overload
  667. def __getitem__(
  668. self: ARRAY_VAR_OF_LIST_ELEMENT[str], i: int | NumberVar
  669. ) -> StringVar: ...
  670. @overload
  671. def __getitem__(
  672. self: ARRAY_VAR_OF_LIST_ELEMENT[bool], i: int | NumberVar
  673. ) -> BooleanVar: ...
  674. @overload
  675. def __getitem__(
  676. self: ARRAY_VAR_OF_LIST_ELEMENT[List[INNER_ARRAY_VAR]],
  677. i: int | NumberVar,
  678. ) -> ArrayVar[List[INNER_ARRAY_VAR]]: ...
  679. @overload
  680. def __getitem__(
  681. self: ARRAY_VAR_OF_LIST_ELEMENT[Set[INNER_ARRAY_VAR]],
  682. i: int | NumberVar,
  683. ) -> ArrayVar[Set[INNER_ARRAY_VAR]]: ...
  684. @overload
  685. def __getitem__(
  686. self: ARRAY_VAR_OF_LIST_ELEMENT[Tuple[KEY_TYPE, VALUE_TYPE]],
  687. i: int | NumberVar,
  688. ) -> ArrayVar[Tuple[KEY_TYPE, VALUE_TYPE]]: ...
  689. @overload
  690. def __getitem__(
  691. self: ARRAY_VAR_OF_LIST_ELEMENT[Tuple[INNER_ARRAY_VAR, ...]],
  692. i: int | NumberVar,
  693. ) -> ArrayVar[Tuple[INNER_ARRAY_VAR, ...]]: ...
  694. @overload
  695. def __getitem__(
  696. self: ARRAY_VAR_OF_LIST_ELEMENT[Dict[KEY_TYPE, VALUE_TYPE]],
  697. i: int | NumberVar,
  698. ) -> ObjectVar[Dict[KEY_TYPE, VALUE_TYPE]]: ...
  699. @overload
  700. def __getitem__(self, i: int | NumberVar) -> Var: ...
  701. def __getitem__(self, i: Any) -> ArrayVar[ARRAY_VAR_TYPE] | Var:
  702. """Get a slice of the array.
  703. Args:
  704. i: The slice.
  705. Returns:
  706. The array slice operation.
  707. """
  708. if isinstance(i, slice):
  709. return ArraySliceOperation.create(self, i)
  710. if not isinstance(i, (int, NumberVar)) or (
  711. isinstance(i, NumberVar) and i._is_strict_float()
  712. ):
  713. raise_unsupported_operand_types("[]", (type(self), type(i)))
  714. return array_item_operation(self, i)
  715. def length(self) -> NumberVar:
  716. """Get the length of the array.
  717. Returns:
  718. The length of the array.
  719. """
  720. return array_length_operation(self)
  721. @overload
  722. @classmethod
  723. def range(cls, stop: int | NumberVar, /) -> ArrayVar[List[int]]: ...
  724. @overload
  725. @classmethod
  726. def range(
  727. cls,
  728. start: int | NumberVar,
  729. end: int | NumberVar,
  730. step: int | NumberVar = 1,
  731. /,
  732. ) -> ArrayVar[List[int]]: ...
  733. @overload
  734. @classmethod
  735. def range(
  736. cls,
  737. first_endpoint: int | NumberVar,
  738. second_endpoint: int | NumberVar | None = None,
  739. step: int | NumberVar | None = None,
  740. ) -> ArrayVar[List[int]]: ...
  741. @classmethod
  742. def range(
  743. cls,
  744. first_endpoint: int | NumberVar,
  745. second_endpoint: int | NumberVar | None = None,
  746. step: int | NumberVar | None = None,
  747. ) -> ArrayVar[List[int]]:
  748. """Create a range of numbers.
  749. Args:
  750. first_endpoint: The end of the range if second_endpoint is not provided, otherwise the start of the range.
  751. second_endpoint: The end of the range.
  752. step: The step of the range.
  753. Returns:
  754. The range of numbers.
  755. """
  756. if any(
  757. not isinstance(i, (int, NumberVar))
  758. for i in (first_endpoint, second_endpoint, step)
  759. if i is not None
  760. ):
  761. raise_unsupported_operand_types(
  762. "range", (type(first_endpoint), type(second_endpoint), type(step))
  763. )
  764. if second_endpoint is None:
  765. start = 0
  766. end = first_endpoint
  767. else:
  768. start = first_endpoint
  769. end = second_endpoint
  770. return array_range_operation(start, end, step or 1)
  771. @overload
  772. def contains(self, other: Any) -> BooleanVar: ...
  773. @overload
  774. def contains(self, other: Any, field: StringVar | str) -> BooleanVar: ...
  775. def contains(self, other: Any, field: Any = None) -> BooleanVar:
  776. """Check if the array contains an element.
  777. Args:
  778. other: The element to check for.
  779. field: The field to check.
  780. Returns:
  781. The array contains operation.
  782. """
  783. if field is not None:
  784. if not isinstance(field, (StringVar, str)):
  785. raise_unsupported_operand_types("contains", (type(self), type(field)))
  786. return array_contains_field_operation(self, other, field)
  787. return array_contains_operation(self, other)
  788. def pluck(self, field: StringVar | str) -> ArrayVar:
  789. """Pluck a field from the array.
  790. Args:
  791. field: The field to pluck from the array.
  792. Returns:
  793. The array pluck operation.
  794. """
  795. return array_pluck_operation(self, field)
  796. @overload
  797. def __mul__(self, other: NumberVar | int) -> ArrayVar[ARRAY_VAR_TYPE]: ...
  798. @overload
  799. def __mul__(self, other: NoReturn) -> NoReturn: ...
  800. def __mul__(self, other: Any) -> ArrayVar[ARRAY_VAR_TYPE]:
  801. """Multiply the sequence by a number or integer.
  802. Parameters:
  803. other: The number or integer to multiply the sequence by.
  804. Returns:
  805. ArrayVar[ARRAY_VAR_TYPE]: The result of multiplying the sequence by the given number or integer.
  806. """
  807. if not isinstance(other, (NumberVar, int)) or (
  808. isinstance(other, NumberVar) and other._is_strict_float()
  809. ):
  810. raise_unsupported_operand_types("*", (type(self), type(other)))
  811. return repeat_array_operation(self, other)
  812. __rmul__ = __mul__ # type: ignore
  813. @overload
  814. def __lt__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  815. @overload
  816. def __lt__(self, other: list | tuple) -> BooleanVar: ...
  817. def __lt__(self, other: Any):
  818. """Check if the array is less than another array.
  819. Args:
  820. other: The other array.
  821. Returns:
  822. The array less than operation.
  823. """
  824. if not isinstance(other, (ArrayVar, list, tuple)):
  825. raise_unsupported_operand_types("<", (type(self), type(other)))
  826. return array_lt_operation(self, other)
  827. @overload
  828. def __gt__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  829. @overload
  830. def __gt__(self, other: list | tuple) -> BooleanVar: ...
  831. def __gt__(self, other: Any):
  832. """Check if the array is greater than another array.
  833. Args:
  834. other: The other array.
  835. Returns:
  836. The array greater than operation.
  837. """
  838. if not isinstance(other, (ArrayVar, list, tuple)):
  839. raise_unsupported_operand_types(">", (type(self), type(other)))
  840. return array_gt_operation(self, other)
  841. @overload
  842. def __le__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  843. @overload
  844. def __le__(self, other: list | tuple) -> BooleanVar: ...
  845. def __le__(self, other: Any):
  846. """Check if the array is less than or equal to another array.
  847. Args:
  848. other: The other array.
  849. Returns:
  850. The array less than or equal operation.
  851. """
  852. if not isinstance(other, (ArrayVar, list, tuple)):
  853. raise_unsupported_operand_types("<=", (type(self), type(other)))
  854. return array_le_operation(self, other)
  855. @overload
  856. def __ge__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  857. @overload
  858. def __ge__(self, other: list | tuple) -> BooleanVar: ...
  859. def __ge__(self, other: Any):
  860. """Check if the array is greater than or equal to another array.
  861. Args:
  862. other: The other array.
  863. Returns:
  864. The array greater than or equal operation.
  865. """
  866. if not isinstance(other, (ArrayVar, list, tuple)):
  867. raise_unsupported_operand_types(">=", (type(self), type(other)))
  868. return array_ge_operation(self, other)
  869. def foreach(self, fn: Any):
  870. """Apply a function to each element of the array.
  871. Args:
  872. fn: The function to apply.
  873. Returns:
  874. The array after applying the function.
  875. Raises:
  876. VarTypeError: If the function takes more than one argument.
  877. """
  878. from .function import ArgsFunctionOperation
  879. if not callable(fn):
  880. raise_unsupported_operand_types("foreach", (type(self), type(fn)))
  881. # get the number of arguments of the function
  882. num_args = len(inspect.signature(fn).parameters)
  883. if num_args > 1:
  884. raise VarTypeError(
  885. "The function passed to foreach should take at most one argument."
  886. )
  887. if num_args == 0:
  888. return_value = fn()
  889. function_var = ArgsFunctionOperation.create(tuple(), return_value)
  890. else:
  891. # generic number var
  892. number_var = Var("").to(NumberVar)
  893. first_arg_type = self[number_var]._var_type
  894. arg_name = get_unique_variable_name()
  895. # get first argument type
  896. first_arg = Var(
  897. _js_expr=arg_name,
  898. _var_type=first_arg_type,
  899. ).guess_type()
  900. function_var = ArgsFunctionOperation.create((arg_name,), fn(first_arg))
  901. return map_array_operation(self, function_var)
  902. LIST_ELEMENT = TypeVar("LIST_ELEMENT")
  903. ARRAY_VAR_OF_LIST_ELEMENT = Union[
  904. ArrayVar[List[LIST_ELEMENT]],
  905. ArrayVar[Set[LIST_ELEMENT]],
  906. ArrayVar[Tuple[LIST_ELEMENT, ...]],
  907. ]
  908. @dataclasses.dataclass(
  909. eq=False,
  910. frozen=True,
  911. **{"slots": True} if sys.version_info >= (3, 10) else {},
  912. )
  913. class LiteralArrayVar(CachedVarOperation, LiteralVar, ArrayVar[ARRAY_VAR_TYPE]):
  914. """Base class for immutable literal array vars."""
  915. _var_value: Union[
  916. List[Union[Var, Any]],
  917. Set[Union[Var, Any]],
  918. Tuple[Union[Var, Any], ...],
  919. ] = dataclasses.field(default_factory=list)
  920. @cached_property_no_lock
  921. def _cached_var_name(self) -> str:
  922. """The name of the var.
  923. Returns:
  924. The name of the var.
  925. """
  926. return (
  927. "["
  928. + ", ".join(
  929. [str(LiteralVar.create(element)) for element in self._var_value]
  930. )
  931. + "]"
  932. )
  933. @cached_property_no_lock
  934. def _cached_get_all_var_data(self) -> VarData | None:
  935. """Get all the VarData associated with the Var.
  936. Returns:
  937. The VarData associated with the Var.
  938. """
  939. return VarData.merge(
  940. *[
  941. LiteralVar.create(element)._get_all_var_data()
  942. for element in self._var_value
  943. ],
  944. self._var_data,
  945. )
  946. def __hash__(self) -> int:
  947. """Get the hash of the var.
  948. Returns:
  949. The hash of the var.
  950. """
  951. return hash((self.__class__.__name__, self._js_expr))
  952. def json(self) -> str:
  953. """Get the JSON representation of the var.
  954. Returns:
  955. The JSON representation of the var.
  956. """
  957. return (
  958. "["
  959. + ", ".join(
  960. [LiteralVar.create(element).json() for element in self._var_value]
  961. )
  962. + "]"
  963. )
  964. @classmethod
  965. def create(
  966. cls,
  967. value: ARRAY_VAR_TYPE,
  968. _var_type: Type[ARRAY_VAR_TYPE] | None = None,
  969. _var_data: VarData | None = None,
  970. ) -> LiteralArrayVar[ARRAY_VAR_TYPE]:
  971. """Create a var from a string value.
  972. Args:
  973. value: The value to create the var from.
  974. _var_data: Additional hooks and imports associated with the Var.
  975. Returns:
  976. The var.
  977. """
  978. return cls(
  979. _js_expr="",
  980. _var_type=figure_out_type(value) if _var_type is None else _var_type,
  981. _var_data=_var_data,
  982. _var_value=value,
  983. )
  984. @var_operation
  985. def string_split_operation(string: StringVar, sep: StringVar | str = ""):
  986. """Split a string.
  987. Args:
  988. string: The string to split.
  989. sep: The separator.
  990. Returns:
  991. The split string.
  992. """
  993. return var_operation_return(
  994. js_expression=f"{string}.split({sep})", var_type=List[str]
  995. )
  996. @dataclasses.dataclass(
  997. eq=False,
  998. frozen=True,
  999. **{"slots": True} if sys.version_info >= (3, 10) else {},
  1000. )
  1001. class ArraySliceOperation(CachedVarOperation, ArrayVar):
  1002. """Base class for immutable string vars that are the result of a string slice operation."""
  1003. _array: ArrayVar = dataclasses.field(
  1004. default_factory=lambda: LiteralArrayVar.create([])
  1005. )
  1006. _start: NumberVar | int = dataclasses.field(default_factory=lambda: 0)
  1007. _stop: NumberVar | int = dataclasses.field(default_factory=lambda: 0)
  1008. _step: NumberVar | int = dataclasses.field(default_factory=lambda: 1)
  1009. @cached_property_no_lock
  1010. def _cached_var_name(self) -> str:
  1011. """The name of the var.
  1012. Returns:
  1013. The name of the var.
  1014. Raises:
  1015. ValueError: If the slice step is zero.
  1016. """
  1017. start, end, step = self._start, self._stop, self._step
  1018. normalized_start = (
  1019. LiteralVar.create(start) if start is not None else Var(_js_expr="undefined")
  1020. )
  1021. normalized_end = (
  1022. LiteralVar.create(end) if end is not None else Var(_js_expr="undefined")
  1023. )
  1024. if step is None:
  1025. return f"{str(self._array)}.slice({str(normalized_start)}, {str(normalized_end)})"
  1026. if not isinstance(step, Var):
  1027. if step < 0:
  1028. actual_start = end + 1 if end is not None else 0
  1029. actual_end = start + 1 if start is not None else self._array.length()
  1030. return str(self._array[actual_start:actual_end].reverse()[::-step])
  1031. if step == 0:
  1032. raise ValueError("slice step cannot be zero")
  1033. return f"{str(self._array)}.slice({str(normalized_start)}, {str(normalized_end)}).filter((_, i) => i % {str(step)} === 0)"
  1034. actual_start_reverse = end + 1 if end is not None else 0
  1035. actual_end_reverse = start + 1 if start is not None else self._array.length()
  1036. return f"{str(self.step)} > 0 ? {str(self._array)}.slice({str(normalized_start)}, {str(normalized_end)}).filter((_, i) => i % {str(step)} === 0) : {str(self._array)}.slice({str(actual_start_reverse)}, {str(actual_end_reverse)}).reverse().filter((_, i) => i % {str(-step)} === 0)"
  1037. @classmethod
  1038. def create(
  1039. cls,
  1040. array: ArrayVar,
  1041. slice: slice,
  1042. _var_data: VarData | None = None,
  1043. ) -> ArraySliceOperation:
  1044. """Create a var from a string value.
  1045. Args:
  1046. array: The array.
  1047. slice: The slice.
  1048. _var_data: Additional hooks and imports associated with the Var.
  1049. Returns:
  1050. The var.
  1051. """
  1052. return cls(
  1053. _js_expr="",
  1054. _var_type=array._var_type,
  1055. _var_data=_var_data,
  1056. _array=array,
  1057. _start=slice.start,
  1058. _stop=slice.stop,
  1059. _step=slice.step,
  1060. )
  1061. @var_operation
  1062. def array_pluck_operation(
  1063. array: ArrayVar[ARRAY_VAR_TYPE],
  1064. field: StringVar | str,
  1065. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1066. """Pluck a field from an array of objects.
  1067. Args:
  1068. array: The array to pluck from.
  1069. field: The field to pluck from the objects in the array.
  1070. Returns:
  1071. The reversed array.
  1072. """
  1073. return var_operation_return(
  1074. js_expression=f"{array}.map(e=>e?.[{field}])",
  1075. var_type=array._var_type,
  1076. )
  1077. @var_operation
  1078. def array_reverse_operation(
  1079. array: ArrayVar[ARRAY_VAR_TYPE],
  1080. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1081. """Reverse an array.
  1082. Args:
  1083. array: The array to reverse.
  1084. Returns:
  1085. The reversed array.
  1086. """
  1087. return var_operation_return(
  1088. js_expression=f"{array}.slice().reverse()",
  1089. var_type=array._var_type,
  1090. )
  1091. @var_operation
  1092. def array_lt_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1093. """Check if an array is less than another array.
  1094. Args:
  1095. lhs: The left-hand side array.
  1096. rhs: The right-hand side array.
  1097. Returns:
  1098. The array less than operation.
  1099. """
  1100. return var_operation_return(js_expression=f"{lhs} < {rhs}", var_type=bool)
  1101. @var_operation
  1102. def array_gt_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1103. """Check if an array is greater than another array.
  1104. Args:
  1105. lhs: The left-hand side array.
  1106. rhs: The right-hand side array.
  1107. Returns:
  1108. The array greater than operation.
  1109. """
  1110. return var_operation_return(js_expression=f"{lhs} > {rhs}", var_type=bool)
  1111. @var_operation
  1112. def array_le_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1113. """Check if an array is less than or equal to another array.
  1114. Args:
  1115. lhs: The left-hand side array.
  1116. rhs: The right-hand side array.
  1117. Returns:
  1118. The array less than or equal operation.
  1119. """
  1120. return var_operation_return(js_expression=f"{lhs} <= {rhs}", var_type=bool)
  1121. @var_operation
  1122. def array_ge_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1123. """Check if an array is greater than or equal to another array.
  1124. Args:
  1125. lhs: The left-hand side array.
  1126. rhs: The right-hand side array.
  1127. Returns:
  1128. The array greater than or equal operation.
  1129. """
  1130. return var_operation_return(js_expression=f"{lhs} >= {rhs}", var_type=bool)
  1131. @var_operation
  1132. def array_length_operation(array: ArrayVar):
  1133. """Get the length of an array.
  1134. Args:
  1135. array: The array.
  1136. Returns:
  1137. The length of the array.
  1138. """
  1139. return var_operation_return(
  1140. js_expression=f"{array}.length",
  1141. var_type=int,
  1142. )
  1143. def is_tuple_type(t: GenericType) -> bool:
  1144. """Check if a type is a tuple type.
  1145. Args:
  1146. t: The type to check.
  1147. Returns:
  1148. Whether the type is a tuple type.
  1149. """
  1150. if inspect.isclass(t):
  1151. return issubclass(t, tuple)
  1152. return get_origin(t) is tuple
  1153. @var_operation
  1154. def array_item_operation(array: ArrayVar, index: NumberVar | int):
  1155. """Get an item from an array.
  1156. Args:
  1157. array: The array.
  1158. index: The index of the item.
  1159. Returns:
  1160. The item from the array.
  1161. """
  1162. args = typing.get_args(array._var_type)
  1163. if args and isinstance(index, LiteralNumberVar) and is_tuple_type(array._var_type):
  1164. index_value = int(index._var_value)
  1165. element_type = args[index_value % len(args)]
  1166. else:
  1167. element_type = unionize(*args)
  1168. return var_operation_return(
  1169. js_expression=f"{str(array)}.at({str(index)})",
  1170. var_type=element_type,
  1171. )
  1172. @var_operation
  1173. def array_range_operation(
  1174. start: NumberVar | int, stop: NumberVar | int, step: NumberVar | int
  1175. ):
  1176. """Create a range of numbers.
  1177. Args:
  1178. start: The start of the range.
  1179. stop: The end of the range.
  1180. step: The step of the range.
  1181. Returns:
  1182. The range of numbers.
  1183. """
  1184. return var_operation_return(
  1185. js_expression=f"Array.from({{ length: ({str(stop)} - {str(start)}) / {str(step)} }}, (_, i) => {str(start)} + i * {str(step)})",
  1186. var_type=List[int],
  1187. )
  1188. @var_operation
  1189. def array_contains_field_operation(
  1190. haystack: ArrayVar, needle: Any | Var, field: StringVar | str
  1191. ):
  1192. """Check if an array contains an element.
  1193. Args:
  1194. haystack: The array to check.
  1195. needle: The element to check for.
  1196. field: The field to check.
  1197. Returns:
  1198. The array contains operation.
  1199. """
  1200. return var_operation_return(
  1201. js_expression=f"{haystack}.some(obj => obj[{field}] === {needle})",
  1202. var_type=bool,
  1203. )
  1204. @var_operation
  1205. def array_contains_operation(haystack: ArrayVar, needle: Any | Var):
  1206. """Check if an array contains an element.
  1207. Args:
  1208. haystack: The array to check.
  1209. needle: The element to check for.
  1210. Returns:
  1211. The array contains operation.
  1212. """
  1213. return var_operation_return(
  1214. js_expression=f"{haystack}.includes({needle})",
  1215. var_type=bool,
  1216. )
  1217. @dataclasses.dataclass(
  1218. eq=False,
  1219. frozen=True,
  1220. **{"slots": True} if sys.version_info >= (3, 10) else {},
  1221. )
  1222. class ToStringOperation(ToOperation, StringVar):
  1223. """Base class for immutable string vars that are the result of a to string operation."""
  1224. _original: Var = dataclasses.field(default_factory=lambda: LiteralNoneVar.create())
  1225. _default_var_type: ClassVar[Type] = str
  1226. @dataclasses.dataclass(
  1227. eq=False,
  1228. frozen=True,
  1229. **{"slots": True} if sys.version_info >= (3, 10) else {},
  1230. )
  1231. class ToArrayOperation(ToOperation, ArrayVar):
  1232. """Base class for immutable array vars that are the result of a to array operation."""
  1233. _original: Var = dataclasses.field(default_factory=lambda: LiteralNoneVar.create())
  1234. _default_var_type: ClassVar[Type] = List[Any]
  1235. @var_operation
  1236. def repeat_array_operation(
  1237. array: ArrayVar[ARRAY_VAR_TYPE], count: NumberVar | int
  1238. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1239. """Repeat an array a number of times.
  1240. Args:
  1241. array: The array to repeat.
  1242. count: The number of times to repeat the array.
  1243. Returns:
  1244. The repeated array.
  1245. """
  1246. return var_operation_return(
  1247. js_expression=f"Array.from({{ length: {count} }}).flatMap(() => {array})",
  1248. var_type=array._var_type,
  1249. )
  1250. if TYPE_CHECKING:
  1251. from .function import FunctionVar
  1252. @var_operation
  1253. def map_array_operation(
  1254. array: ArrayVar[ARRAY_VAR_TYPE],
  1255. function: FunctionVar,
  1256. ):
  1257. """Map a function over an array.
  1258. Args:
  1259. array: The array.
  1260. function: The function to map.
  1261. Returns:
  1262. The mapped array.
  1263. """
  1264. return var_operation_return(
  1265. js_expression=f"{array}.map({function})", var_type=List[Any]
  1266. )
  1267. @var_operation
  1268. def array_concat_operation(
  1269. lhs: ArrayVar[ARRAY_VAR_TYPE], rhs: ArrayVar[ARRAY_VAR_TYPE]
  1270. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1271. """Concatenate two arrays.
  1272. Args:
  1273. lhs: The left-hand side array.
  1274. rhs: The right-hand side array.
  1275. Returns:
  1276. The concatenated array.
  1277. """
  1278. return var_operation_return(
  1279. js_expression=f"[...{lhs}, ...{rhs}]",
  1280. var_type=Union[lhs._var_type, rhs._var_type],
  1281. )