sequence.py 45 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656
  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 = None,
  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. # Determine var type in case the value is inherited from str.
  421. _var_type = _var_type or type(value) or str
  422. if REFLEX_VAR_OPENING_TAG in value:
  423. strings_and_vals: list[Var | str] = []
  424. offset = 0
  425. # Find all tags
  426. while m := _decode_var_pattern.search(value):
  427. start, end = m.span()
  428. strings_and_vals.append(value[:start])
  429. serialized_data = m.group(1)
  430. if serialized_data.isnumeric() or (
  431. serialized_data[0] == "-" and serialized_data[1:].isnumeric()
  432. ):
  433. # This is a global immutable var.
  434. var = _global_vars[int(serialized_data)]
  435. strings_and_vals.append(var)
  436. value = value[(end + len(var._js_expr)) :]
  437. offset += end - start
  438. strings_and_vals.append(value)
  439. filtered_strings_and_vals = [
  440. s for s in strings_and_vals if isinstance(s, Var) or s
  441. ]
  442. if len(filtered_strings_and_vals) == 1:
  443. only_string = filtered_strings_and_vals[0]
  444. if isinstance(only_string, str):
  445. return LiteralVar.create(only_string).to(StringVar, _var_type)
  446. else:
  447. return only_string.to(StringVar, only_string._var_type)
  448. if len(
  449. literal_strings := [
  450. s
  451. for s in filtered_strings_and_vals
  452. if isinstance(s, (str, LiteralStringVar))
  453. ]
  454. ) == len(filtered_strings_and_vals):
  455. return LiteralStringVar.create(
  456. "".join(
  457. s._var_value if isinstance(s, LiteralStringVar) else s
  458. for s in literal_strings
  459. ),
  460. _var_type=_var_type,
  461. _var_data=VarData.merge(
  462. _var_data,
  463. *(
  464. s._get_all_var_data()
  465. for s in filtered_strings_and_vals
  466. if isinstance(s, Var)
  467. ),
  468. ),
  469. )
  470. concat_result = ConcatVarOperation.create(
  471. *filtered_strings_and_vals,
  472. _var_data=_var_data,
  473. )
  474. return (
  475. concat_result
  476. if _var_type is str
  477. else concat_result.to(StringVar, _var_type)
  478. )
  479. return LiteralStringVar(
  480. _js_expr=json.dumps(value),
  481. _var_type=_var_type,
  482. _var_data=_var_data,
  483. _var_value=value,
  484. )
  485. def __hash__(self) -> int:
  486. """Get the hash of the var.
  487. Returns:
  488. The hash of the var.
  489. """
  490. return hash((self.__class__.__name__, self._var_value))
  491. def json(self) -> str:
  492. """Get the JSON representation of the var.
  493. Returns:
  494. The JSON representation of the var.
  495. """
  496. return json.dumps(self._var_value)
  497. @dataclasses.dataclass(
  498. eq=False,
  499. frozen=True,
  500. **{"slots": True} if sys.version_info >= (3, 10) else {},
  501. )
  502. class ConcatVarOperation(CachedVarOperation, StringVar):
  503. """Representing a concatenation of literal string vars."""
  504. _var_value: Tuple[Var, ...] = dataclasses.field(default_factory=tuple)
  505. @cached_property_no_lock
  506. def _cached_var_name(self) -> str:
  507. """The name of the var.
  508. Returns:
  509. The name of the var.
  510. """
  511. list_of_strs: List[Union[str, Var]] = []
  512. last_string = ""
  513. for var in self._var_value:
  514. if isinstance(var, LiteralStringVar):
  515. last_string += var._var_value
  516. else:
  517. if last_string:
  518. list_of_strs.append(last_string)
  519. last_string = ""
  520. list_of_strs.append(var)
  521. if last_string:
  522. list_of_strs.append(last_string)
  523. list_of_strs_filtered = [
  524. str(LiteralVar.create(s)) for s in list_of_strs if isinstance(s, Var) or s
  525. ]
  526. if len(list_of_strs_filtered) == 1:
  527. return list_of_strs_filtered[0]
  528. return "(" + "+".join(list_of_strs_filtered) + ")"
  529. @cached_property_no_lock
  530. def _cached_get_all_var_data(self) -> VarData | None:
  531. """Get all the VarData asVarDatae Var.
  532. Returns:
  533. The VarData associated with the Var.
  534. """
  535. return VarData.merge(
  536. *[
  537. var._get_all_var_data()
  538. for var in self._var_value
  539. if isinstance(var, Var)
  540. ],
  541. self._var_data,
  542. )
  543. @classmethod
  544. def create(
  545. cls,
  546. *value: Var | str,
  547. _var_data: VarData | None = None,
  548. ) -> ConcatVarOperation:
  549. """Create a var from a string value.
  550. Args:
  551. value: The values to concatenate.
  552. _var_data: Additional hooks and imports associated with the Var.
  553. Returns:
  554. The var.
  555. """
  556. return cls(
  557. _js_expr="",
  558. _var_type=str,
  559. _var_data=_var_data,
  560. _var_value=tuple(map(LiteralVar.create, value)),
  561. )
  562. ARRAY_VAR_TYPE = TypeVar("ARRAY_VAR_TYPE", bound=Union[List, Tuple, Set])
  563. OTHER_TUPLE = TypeVar("OTHER_TUPLE")
  564. INNER_ARRAY_VAR = TypeVar("INNER_ARRAY_VAR")
  565. KEY_TYPE = TypeVar("KEY_TYPE")
  566. VALUE_TYPE = TypeVar("VALUE_TYPE")
  567. class ArrayVar(Var[ARRAY_VAR_TYPE]):
  568. """Base class for immutable array vars."""
  569. @overload
  570. def join(self, sep: StringVar | str = "") -> StringVar: ...
  571. @overload
  572. def join(self, sep: NoReturn) -> NoReturn: ...
  573. def join(self, sep: Any = "") -> StringVar:
  574. """Join the elements of the array.
  575. Args:
  576. sep: The separator between elements.
  577. Returns:
  578. The joined elements.
  579. """
  580. if not isinstance(sep, (StringVar, str)):
  581. raise_unsupported_operand_types("join", (type(self), type(sep)))
  582. if (
  583. isinstance(self, LiteralArrayVar)
  584. and (
  585. len(
  586. args := [
  587. x
  588. for x in self._var_value
  589. if isinstance(x, (LiteralStringVar, str))
  590. ]
  591. )
  592. == len(self._var_value)
  593. )
  594. and isinstance(sep, (LiteralStringVar, str))
  595. ):
  596. sep_str = sep._var_value if isinstance(sep, LiteralStringVar) else sep
  597. return LiteralStringVar.create(
  598. sep_str.join(
  599. i._var_value if isinstance(i, LiteralStringVar) else i for i in args
  600. )
  601. )
  602. return array_join_operation(self, sep)
  603. def reverse(self) -> ArrayVar[ARRAY_VAR_TYPE]:
  604. """Reverse the array.
  605. Returns:
  606. The reversed array.
  607. """
  608. return array_reverse_operation(self)
  609. @overload
  610. def __add__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> ArrayVar[ARRAY_VAR_TYPE]: ...
  611. @overload
  612. def __add__(self, other: NoReturn) -> NoReturn: ...
  613. def __add__(self, other: Any) -> ArrayVar[ARRAY_VAR_TYPE]:
  614. """Concatenate two arrays.
  615. Parameters:
  616. other: The other array to concatenate.
  617. Returns:
  618. ArrayConcatOperation: The concatenation of the two arrays.
  619. """
  620. if not isinstance(other, ArrayVar):
  621. raise_unsupported_operand_types("+", (type(self), type(other)))
  622. return array_concat_operation(self, other)
  623. @overload
  624. def __getitem__(self, i: slice) -> ArrayVar[ARRAY_VAR_TYPE]: ...
  625. @overload
  626. def __getitem__(
  627. self: (
  628. ArrayVar[Tuple[int, OTHER_TUPLE]]
  629. | ArrayVar[Tuple[float, OTHER_TUPLE]]
  630. | ArrayVar[Tuple[int | float, OTHER_TUPLE]]
  631. ),
  632. i: Literal[0, -2],
  633. ) -> NumberVar: ...
  634. @overload
  635. def __getitem__(
  636. self: (
  637. ArrayVar[Tuple[OTHER_TUPLE, int]]
  638. | ArrayVar[Tuple[OTHER_TUPLE, float]]
  639. | ArrayVar[Tuple[OTHER_TUPLE, int | float]]
  640. ),
  641. i: Literal[1, -1],
  642. ) -> NumberVar: ...
  643. @overload
  644. def __getitem__(
  645. self: ArrayVar[Tuple[str, OTHER_TUPLE]], i: Literal[0, -2]
  646. ) -> StringVar: ...
  647. @overload
  648. def __getitem__(
  649. self: ArrayVar[Tuple[OTHER_TUPLE, str]], i: Literal[1, -1]
  650. ) -> StringVar: ...
  651. @overload
  652. def __getitem__(
  653. self: ArrayVar[Tuple[bool, OTHER_TUPLE]], i: Literal[0, -2]
  654. ) -> BooleanVar: ...
  655. @overload
  656. def __getitem__(
  657. self: ArrayVar[Tuple[OTHER_TUPLE, bool]], i: Literal[1, -1]
  658. ) -> BooleanVar: ...
  659. @overload
  660. def __getitem__(
  661. self: (
  662. ARRAY_VAR_OF_LIST_ELEMENT[int]
  663. | ARRAY_VAR_OF_LIST_ELEMENT[float]
  664. | ARRAY_VAR_OF_LIST_ELEMENT[int | float]
  665. ),
  666. i: int | NumberVar,
  667. ) -> NumberVar: ...
  668. @overload
  669. def __getitem__(
  670. self: ARRAY_VAR_OF_LIST_ELEMENT[str], i: int | NumberVar
  671. ) -> StringVar: ...
  672. @overload
  673. def __getitem__(
  674. self: ARRAY_VAR_OF_LIST_ELEMENT[bool], i: int | NumberVar
  675. ) -> BooleanVar: ...
  676. @overload
  677. def __getitem__(
  678. self: ARRAY_VAR_OF_LIST_ELEMENT[List[INNER_ARRAY_VAR]],
  679. i: int | NumberVar,
  680. ) -> ArrayVar[List[INNER_ARRAY_VAR]]: ...
  681. @overload
  682. def __getitem__(
  683. self: ARRAY_VAR_OF_LIST_ELEMENT[Set[INNER_ARRAY_VAR]],
  684. i: int | NumberVar,
  685. ) -> ArrayVar[Set[INNER_ARRAY_VAR]]: ...
  686. @overload
  687. def __getitem__(
  688. self: ARRAY_VAR_OF_LIST_ELEMENT[Tuple[KEY_TYPE, VALUE_TYPE]],
  689. i: int | NumberVar,
  690. ) -> ArrayVar[Tuple[KEY_TYPE, VALUE_TYPE]]: ...
  691. @overload
  692. def __getitem__(
  693. self: ARRAY_VAR_OF_LIST_ELEMENT[Tuple[INNER_ARRAY_VAR, ...]],
  694. i: int | NumberVar,
  695. ) -> ArrayVar[Tuple[INNER_ARRAY_VAR, ...]]: ...
  696. @overload
  697. def __getitem__(
  698. self: ARRAY_VAR_OF_LIST_ELEMENT[Dict[KEY_TYPE, VALUE_TYPE]],
  699. i: int | NumberVar,
  700. ) -> ObjectVar[Dict[KEY_TYPE, VALUE_TYPE]]: ...
  701. @overload
  702. def __getitem__(self, i: int | NumberVar) -> Var: ...
  703. def __getitem__(self, i: Any) -> ArrayVar[ARRAY_VAR_TYPE] | Var:
  704. """Get a slice of the array.
  705. Args:
  706. i: The slice.
  707. Returns:
  708. The array slice operation.
  709. """
  710. if isinstance(i, slice):
  711. return ArraySliceOperation.create(self, i)
  712. if not isinstance(i, (int, NumberVar)) or (
  713. isinstance(i, NumberVar) and i._is_strict_float()
  714. ):
  715. raise_unsupported_operand_types("[]", (type(self), type(i)))
  716. return array_item_operation(self, i)
  717. def length(self) -> NumberVar:
  718. """Get the length of the array.
  719. Returns:
  720. The length of the array.
  721. """
  722. return array_length_operation(self)
  723. @overload
  724. @classmethod
  725. def range(cls, stop: int | NumberVar, /) -> ArrayVar[List[int]]: ...
  726. @overload
  727. @classmethod
  728. def range(
  729. cls,
  730. start: int | NumberVar,
  731. end: int | NumberVar,
  732. step: int | NumberVar = 1,
  733. /,
  734. ) -> ArrayVar[List[int]]: ...
  735. @overload
  736. @classmethod
  737. def range(
  738. cls,
  739. first_endpoint: int | NumberVar,
  740. second_endpoint: int | NumberVar | None = None,
  741. step: int | NumberVar | None = None,
  742. ) -> ArrayVar[List[int]]: ...
  743. @classmethod
  744. def range(
  745. cls,
  746. first_endpoint: int | NumberVar,
  747. second_endpoint: int | NumberVar | None = None,
  748. step: int | NumberVar | None = None,
  749. ) -> ArrayVar[List[int]]:
  750. """Create a range of numbers.
  751. Args:
  752. first_endpoint: The end of the range if second_endpoint is not provided, otherwise the start of the range.
  753. second_endpoint: The end of the range.
  754. step: The step of the range.
  755. Returns:
  756. The range of numbers.
  757. """
  758. if any(
  759. not isinstance(i, (int, NumberVar))
  760. for i in (first_endpoint, second_endpoint, step)
  761. if i is not None
  762. ):
  763. raise_unsupported_operand_types(
  764. "range", (type(first_endpoint), type(second_endpoint), type(step))
  765. )
  766. if second_endpoint is None:
  767. start = 0
  768. end = first_endpoint
  769. else:
  770. start = first_endpoint
  771. end = second_endpoint
  772. return array_range_operation(start, end, step or 1)
  773. @overload
  774. def contains(self, other: Any) -> BooleanVar: ...
  775. @overload
  776. def contains(self, other: Any, field: StringVar | str) -> BooleanVar: ...
  777. def contains(self, other: Any, field: Any = None) -> BooleanVar:
  778. """Check if the array contains an element.
  779. Args:
  780. other: The element to check for.
  781. field: The field to check.
  782. Returns:
  783. The array contains operation.
  784. """
  785. if field is not None:
  786. if not isinstance(field, (StringVar, str)):
  787. raise_unsupported_operand_types("contains", (type(self), type(field)))
  788. return array_contains_field_operation(self, other, field)
  789. return array_contains_operation(self, other)
  790. def pluck(self, field: StringVar | str) -> ArrayVar:
  791. """Pluck a field from the array.
  792. Args:
  793. field: The field to pluck from the array.
  794. Returns:
  795. The array pluck operation.
  796. """
  797. return array_pluck_operation(self, field)
  798. @overload
  799. def __mul__(self, other: NumberVar | int) -> ArrayVar[ARRAY_VAR_TYPE]: ...
  800. @overload
  801. def __mul__(self, other: NoReturn) -> NoReturn: ...
  802. def __mul__(self, other: Any) -> ArrayVar[ARRAY_VAR_TYPE]:
  803. """Multiply the sequence by a number or integer.
  804. Parameters:
  805. other: The number or integer to multiply the sequence by.
  806. Returns:
  807. ArrayVar[ARRAY_VAR_TYPE]: The result of multiplying the sequence by the given number or integer.
  808. """
  809. if not isinstance(other, (NumberVar, int)) or (
  810. isinstance(other, NumberVar) and other._is_strict_float()
  811. ):
  812. raise_unsupported_operand_types("*", (type(self), type(other)))
  813. return repeat_array_operation(self, other)
  814. __rmul__ = __mul__ # type: ignore
  815. @overload
  816. def __lt__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  817. @overload
  818. def __lt__(self, other: list | tuple) -> BooleanVar: ...
  819. def __lt__(self, other: Any):
  820. """Check if the array is less than another array.
  821. Args:
  822. other: The other array.
  823. Returns:
  824. The array less than operation.
  825. """
  826. if not isinstance(other, (ArrayVar, list, tuple)):
  827. raise_unsupported_operand_types("<", (type(self), type(other)))
  828. return array_lt_operation(self, other)
  829. @overload
  830. def __gt__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  831. @overload
  832. def __gt__(self, other: list | tuple) -> BooleanVar: ...
  833. def __gt__(self, other: Any):
  834. """Check if the array is greater than another array.
  835. Args:
  836. other: The other array.
  837. Returns:
  838. The array greater than operation.
  839. """
  840. if not isinstance(other, (ArrayVar, list, tuple)):
  841. raise_unsupported_operand_types(">", (type(self), type(other)))
  842. return array_gt_operation(self, other)
  843. @overload
  844. def __le__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  845. @overload
  846. def __le__(self, other: list | tuple) -> BooleanVar: ...
  847. def __le__(self, other: Any):
  848. """Check if the array is less than or equal to another array.
  849. Args:
  850. other: The other array.
  851. Returns:
  852. The array less than or equal operation.
  853. """
  854. if not isinstance(other, (ArrayVar, list, tuple)):
  855. raise_unsupported_operand_types("<=", (type(self), type(other)))
  856. return array_le_operation(self, other)
  857. @overload
  858. def __ge__(self, other: ArrayVar[ARRAY_VAR_TYPE]) -> BooleanVar: ...
  859. @overload
  860. def __ge__(self, other: list | tuple) -> BooleanVar: ...
  861. def __ge__(self, other: Any):
  862. """Check if the array is greater than or equal to another array.
  863. Args:
  864. other: The other array.
  865. Returns:
  866. The array greater than or equal operation.
  867. """
  868. if not isinstance(other, (ArrayVar, list, tuple)):
  869. raise_unsupported_operand_types(">=", (type(self), type(other)))
  870. return array_ge_operation(self, other)
  871. def foreach(self, fn: Any):
  872. """Apply a function to each element of the array.
  873. Args:
  874. fn: The function to apply.
  875. Returns:
  876. The array after applying the function.
  877. Raises:
  878. VarTypeError: If the function takes more than one argument.
  879. """
  880. from .function import ArgsFunctionOperation
  881. if not callable(fn):
  882. raise_unsupported_operand_types("foreach", (type(self), type(fn)))
  883. # get the number of arguments of the function
  884. num_args = len(inspect.signature(fn).parameters)
  885. if num_args > 1:
  886. raise VarTypeError(
  887. "The function passed to foreach should take at most one argument."
  888. )
  889. if num_args == 0:
  890. return_value = fn()
  891. function_var = ArgsFunctionOperation.create(tuple(), return_value)
  892. else:
  893. # generic number var
  894. number_var = Var("").to(NumberVar)
  895. first_arg_type = self[number_var]._var_type
  896. arg_name = get_unique_variable_name()
  897. # get first argument type
  898. first_arg = Var(
  899. _js_expr=arg_name,
  900. _var_type=first_arg_type,
  901. ).guess_type()
  902. function_var = ArgsFunctionOperation.create((arg_name,), fn(first_arg))
  903. return map_array_operation(self, function_var)
  904. LIST_ELEMENT = TypeVar("LIST_ELEMENT")
  905. ARRAY_VAR_OF_LIST_ELEMENT = Union[
  906. ArrayVar[List[LIST_ELEMENT]],
  907. ArrayVar[Set[LIST_ELEMENT]],
  908. ArrayVar[Tuple[LIST_ELEMENT, ...]],
  909. ]
  910. @dataclasses.dataclass(
  911. eq=False,
  912. frozen=True,
  913. **{"slots": True} if sys.version_info >= (3, 10) else {},
  914. )
  915. class LiteralArrayVar(CachedVarOperation, LiteralVar, ArrayVar[ARRAY_VAR_TYPE]):
  916. """Base class for immutable literal array vars."""
  917. _var_value: Union[
  918. List[Union[Var, Any]],
  919. Set[Union[Var, Any]],
  920. Tuple[Union[Var, Any], ...],
  921. ] = dataclasses.field(default_factory=list)
  922. @cached_property_no_lock
  923. def _cached_var_name(self) -> str:
  924. """The name of the var.
  925. Returns:
  926. The name of the var.
  927. """
  928. return (
  929. "["
  930. + ", ".join(
  931. [str(LiteralVar.create(element)) for element in self._var_value]
  932. )
  933. + "]"
  934. )
  935. @cached_property_no_lock
  936. def _cached_get_all_var_data(self) -> VarData | None:
  937. """Get all the VarData associated with the Var.
  938. Returns:
  939. The VarData associated with the Var.
  940. """
  941. return VarData.merge(
  942. *[
  943. LiteralVar.create(element)._get_all_var_data()
  944. for element in self._var_value
  945. ],
  946. self._var_data,
  947. )
  948. def __hash__(self) -> int:
  949. """Get the hash of the var.
  950. Returns:
  951. The hash of the var.
  952. """
  953. return hash((self.__class__.__name__, self._js_expr))
  954. def json(self) -> str:
  955. """Get the JSON representation of the var.
  956. Returns:
  957. The JSON representation of the var.
  958. """
  959. return (
  960. "["
  961. + ", ".join(
  962. [LiteralVar.create(element).json() for element in self._var_value]
  963. )
  964. + "]"
  965. )
  966. @classmethod
  967. def create(
  968. cls,
  969. value: ARRAY_VAR_TYPE,
  970. _var_type: Type[ARRAY_VAR_TYPE] | None = None,
  971. _var_data: VarData | None = None,
  972. ) -> LiteralArrayVar[ARRAY_VAR_TYPE]:
  973. """Create a var from a string value.
  974. Args:
  975. value: The value to create the var from.
  976. _var_data: Additional hooks and imports associated with the Var.
  977. Returns:
  978. The var.
  979. """
  980. return cls(
  981. _js_expr="",
  982. _var_type=figure_out_type(value) if _var_type is None else _var_type,
  983. _var_data=_var_data,
  984. _var_value=value,
  985. )
  986. @var_operation
  987. def string_split_operation(string: StringVar, sep: StringVar | str = ""):
  988. """Split a string.
  989. Args:
  990. string: The string to split.
  991. sep: The separator.
  992. Returns:
  993. The split string.
  994. """
  995. return var_operation_return(
  996. js_expression=f"{string}.split({sep})", var_type=List[str]
  997. )
  998. @dataclasses.dataclass(
  999. eq=False,
  1000. frozen=True,
  1001. **{"slots": True} if sys.version_info >= (3, 10) else {},
  1002. )
  1003. class ArraySliceOperation(CachedVarOperation, ArrayVar):
  1004. """Base class for immutable string vars that are the result of a string slice operation."""
  1005. _array: ArrayVar = dataclasses.field(
  1006. default_factory=lambda: LiteralArrayVar.create([])
  1007. )
  1008. _start: NumberVar | int = dataclasses.field(default_factory=lambda: 0)
  1009. _stop: NumberVar | int = dataclasses.field(default_factory=lambda: 0)
  1010. _step: NumberVar | int = dataclasses.field(default_factory=lambda: 1)
  1011. @cached_property_no_lock
  1012. def _cached_var_name(self) -> str:
  1013. """The name of the var.
  1014. Returns:
  1015. The name of the var.
  1016. Raises:
  1017. ValueError: If the slice step is zero.
  1018. """
  1019. start, end, step = self._start, self._stop, self._step
  1020. normalized_start = (
  1021. LiteralVar.create(start) if start is not None else Var(_js_expr="undefined")
  1022. )
  1023. normalized_end = (
  1024. LiteralVar.create(end) if end is not None else Var(_js_expr="undefined")
  1025. )
  1026. if step is None:
  1027. return f"{str(self._array)}.slice({str(normalized_start)}, {str(normalized_end)})"
  1028. if not isinstance(step, Var):
  1029. if step < 0:
  1030. actual_start = end + 1 if end is not None else 0
  1031. actual_end = start + 1 if start is not None else self._array.length()
  1032. return str(self._array[actual_start:actual_end].reverse()[::-step])
  1033. if step == 0:
  1034. raise ValueError("slice step cannot be zero")
  1035. return f"{str(self._array)}.slice({str(normalized_start)}, {str(normalized_end)}).filter((_, i) => i % {str(step)} === 0)"
  1036. actual_start_reverse = end + 1 if end is not None else 0
  1037. actual_end_reverse = start + 1 if start is not None else self._array.length()
  1038. 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)"
  1039. @classmethod
  1040. def create(
  1041. cls,
  1042. array: ArrayVar,
  1043. slice: slice,
  1044. _var_data: VarData | None = None,
  1045. ) -> ArraySliceOperation:
  1046. """Create a var from a string value.
  1047. Args:
  1048. array: The array.
  1049. slice: The slice.
  1050. _var_data: Additional hooks and imports associated with the Var.
  1051. Returns:
  1052. The var.
  1053. """
  1054. return cls(
  1055. _js_expr="",
  1056. _var_type=array._var_type,
  1057. _var_data=_var_data,
  1058. _array=array,
  1059. _start=slice.start,
  1060. _stop=slice.stop,
  1061. _step=slice.step,
  1062. )
  1063. @var_operation
  1064. def array_pluck_operation(
  1065. array: ArrayVar[ARRAY_VAR_TYPE],
  1066. field: StringVar | str,
  1067. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1068. """Pluck a field from an array of objects.
  1069. Args:
  1070. array: The array to pluck from.
  1071. field: The field to pluck from the objects in the array.
  1072. Returns:
  1073. The reversed array.
  1074. """
  1075. return var_operation_return(
  1076. js_expression=f"{array}.map(e=>e?.[{field}])",
  1077. var_type=array._var_type,
  1078. )
  1079. @var_operation
  1080. def array_reverse_operation(
  1081. array: ArrayVar[ARRAY_VAR_TYPE],
  1082. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1083. """Reverse an array.
  1084. Args:
  1085. array: The array to reverse.
  1086. Returns:
  1087. The reversed array.
  1088. """
  1089. return var_operation_return(
  1090. js_expression=f"{array}.slice().reverse()",
  1091. var_type=array._var_type,
  1092. )
  1093. @var_operation
  1094. def array_lt_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1095. """Check if an array is less than another array.
  1096. Args:
  1097. lhs: The left-hand side array.
  1098. rhs: The right-hand side array.
  1099. Returns:
  1100. The array less than operation.
  1101. """
  1102. return var_operation_return(js_expression=f"{lhs} < {rhs}", var_type=bool)
  1103. @var_operation
  1104. def array_gt_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1105. """Check if an array is greater than another array.
  1106. Args:
  1107. lhs: The left-hand side array.
  1108. rhs: The right-hand side array.
  1109. Returns:
  1110. The array greater than operation.
  1111. """
  1112. return var_operation_return(js_expression=f"{lhs} > {rhs}", var_type=bool)
  1113. @var_operation
  1114. def array_le_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1115. """Check if an array is less than or equal to another array.
  1116. Args:
  1117. lhs: The left-hand side array.
  1118. rhs: The right-hand side array.
  1119. Returns:
  1120. The array less than or equal operation.
  1121. """
  1122. return var_operation_return(js_expression=f"{lhs} <= {rhs}", var_type=bool)
  1123. @var_operation
  1124. def array_ge_operation(lhs: ArrayVar | list | tuple, rhs: ArrayVar | list | tuple):
  1125. """Check if an array is greater than or equal to another array.
  1126. Args:
  1127. lhs: The left-hand side array.
  1128. rhs: The right-hand side array.
  1129. Returns:
  1130. The array greater than or equal operation.
  1131. """
  1132. return var_operation_return(js_expression=f"{lhs} >= {rhs}", var_type=bool)
  1133. @var_operation
  1134. def array_length_operation(array: ArrayVar):
  1135. """Get the length of an array.
  1136. Args:
  1137. array: The array.
  1138. Returns:
  1139. The length of the array.
  1140. """
  1141. return var_operation_return(
  1142. js_expression=f"{array}.length",
  1143. var_type=int,
  1144. )
  1145. def is_tuple_type(t: GenericType) -> bool:
  1146. """Check if a type is a tuple type.
  1147. Args:
  1148. t: The type to check.
  1149. Returns:
  1150. Whether the type is a tuple type.
  1151. """
  1152. if inspect.isclass(t):
  1153. return issubclass(t, tuple)
  1154. return get_origin(t) is tuple
  1155. @var_operation
  1156. def array_item_operation(array: ArrayVar, index: NumberVar | int):
  1157. """Get an item from an array.
  1158. Args:
  1159. array: The array.
  1160. index: The index of the item.
  1161. Returns:
  1162. The item from the array.
  1163. """
  1164. args = typing.get_args(array._var_type)
  1165. if args and isinstance(index, LiteralNumberVar) and is_tuple_type(array._var_type):
  1166. index_value = int(index._var_value)
  1167. element_type = args[index_value % len(args)]
  1168. else:
  1169. element_type = unionize(*args)
  1170. return var_operation_return(
  1171. js_expression=f"{str(array)}.at({str(index)})",
  1172. var_type=element_type,
  1173. )
  1174. @var_operation
  1175. def array_range_operation(
  1176. start: NumberVar | int, stop: NumberVar | int, step: NumberVar | int
  1177. ):
  1178. """Create a range of numbers.
  1179. Args:
  1180. start: The start of the range.
  1181. stop: The end of the range.
  1182. step: The step of the range.
  1183. Returns:
  1184. The range of numbers.
  1185. """
  1186. return var_operation_return(
  1187. js_expression=f"Array.from({{ length: ({str(stop)} - {str(start)}) / {str(step)} }}, (_, i) => {str(start)} + i * {str(step)})",
  1188. var_type=List[int],
  1189. )
  1190. @var_operation
  1191. def array_contains_field_operation(
  1192. haystack: ArrayVar, needle: Any | Var, field: StringVar | str
  1193. ):
  1194. """Check if an array contains an element.
  1195. Args:
  1196. haystack: The array to check.
  1197. needle: The element to check for.
  1198. field: The field to check.
  1199. Returns:
  1200. The array contains operation.
  1201. """
  1202. return var_operation_return(
  1203. js_expression=f"{haystack}.some(obj => obj[{field}] === {needle})",
  1204. var_type=bool,
  1205. )
  1206. @var_operation
  1207. def array_contains_operation(haystack: ArrayVar, needle: Any | Var):
  1208. """Check if an array contains an element.
  1209. Args:
  1210. haystack: The array to check.
  1211. needle: The element to check for.
  1212. Returns:
  1213. The array contains operation.
  1214. """
  1215. return var_operation_return(
  1216. js_expression=f"{haystack}.includes({needle})",
  1217. var_type=bool,
  1218. )
  1219. @dataclasses.dataclass(
  1220. eq=False,
  1221. frozen=True,
  1222. **{"slots": True} if sys.version_info >= (3, 10) else {},
  1223. )
  1224. class ToStringOperation(ToOperation, StringVar):
  1225. """Base class for immutable string vars that are the result of a to string operation."""
  1226. _original: Var = dataclasses.field(default_factory=lambda: LiteralNoneVar.create())
  1227. _default_var_type: ClassVar[Type] = str
  1228. @dataclasses.dataclass(
  1229. eq=False,
  1230. frozen=True,
  1231. **{"slots": True} if sys.version_info >= (3, 10) else {},
  1232. )
  1233. class ToArrayOperation(ToOperation, ArrayVar):
  1234. """Base class for immutable array vars that are the result of a to array operation."""
  1235. _original: Var = dataclasses.field(default_factory=lambda: LiteralNoneVar.create())
  1236. _default_var_type: ClassVar[Type] = List[Any]
  1237. @var_operation
  1238. def repeat_array_operation(
  1239. array: ArrayVar[ARRAY_VAR_TYPE], count: NumberVar | int
  1240. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1241. """Repeat an array a number of times.
  1242. Args:
  1243. array: The array to repeat.
  1244. count: The number of times to repeat the array.
  1245. Returns:
  1246. The repeated array.
  1247. """
  1248. return var_operation_return(
  1249. js_expression=f"Array.from({{ length: {count} }}).flatMap(() => {array})",
  1250. var_type=array._var_type,
  1251. )
  1252. if TYPE_CHECKING:
  1253. from .function import FunctionVar
  1254. @var_operation
  1255. def map_array_operation(
  1256. array: ArrayVar[ARRAY_VAR_TYPE],
  1257. function: FunctionVar,
  1258. ):
  1259. """Map a function over an array.
  1260. Args:
  1261. array: The array.
  1262. function: The function to map.
  1263. Returns:
  1264. The mapped array.
  1265. """
  1266. return var_operation_return(
  1267. js_expression=f"{array}.map({function})", var_type=List[Any]
  1268. )
  1269. @var_operation
  1270. def array_concat_operation(
  1271. lhs: ArrayVar[ARRAY_VAR_TYPE], rhs: ArrayVar[ARRAY_VAR_TYPE]
  1272. ) -> CustomVarOperationReturn[ARRAY_VAR_TYPE]:
  1273. """Concatenate two arrays.
  1274. Args:
  1275. lhs: The left-hand side array.
  1276. rhs: The right-hand side array.
  1277. Returns:
  1278. The concatenated array.
  1279. """
  1280. return var_operation_return(
  1281. js_expression=f"[...{lhs}, ...{rhs}]",
  1282. var_type=Union[lhs._var_type, rhs._var_type],
  1283. )