pixeltable 0.3.14__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/functions/string.py

@@ -1,5 +1,5 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `StringType`.
+Pixeltable UDFs for `StringType`.
 It closely follows the Pandas `pandas.Series.str` API.
 
 Example:
@@ -12,7 +12,12 @@ t.select(t.str_col.capitalize()).collect()
 """
 
 import builtins
-from typing import Any, Optional
+import re
+import textwrap
+from string import whitespace
+from typing import Any
+
+import sqlalchemy as sql
 
 import pixeltable as pxt
 from pixeltable.utils.code import local_public_names
@@ -28,6 +33,11 @@ def capitalize(self: str) -> str:
     return self.capitalize()
 
 
+@capitalize.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.concat(sql.func.upper(sql.func.left(self, 1)), sql.func.lower(sql.func.right(self, -1)))
+
+
 @pxt.udf(is_method=True)
 def casefold(self: str) -> str:
     """
@@ -53,26 +63,45 @@ def center(self: str, width: int, fillchar: str = ' ') -> str:
 
 
 @pxt.udf(is_method=True)
-def contains(self: str, pattern: str, case: bool = True, flags: int = 0, regex: bool = True) -> bool:
+def contains(self: str, substr: str, case: bool = True) -> bool:
     """
-    Test if string contains pattern or regex.
+    Test if string contains a substring.
 
     Args:
-        pattern: string literal or regular expression
+        substr: string literal or regular expression
         case: if False, ignore case
-        flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
-        regex: if True, treat pattern as a regular expression
     """
-    if regex:
-        import re
+    if case:
+        return substr in self
+    else:
+        return substr.lower() in self.lower()
+
 
-        if not case:
-            flags |= re.IGNORECASE
-        return bool(re.search(pattern, self, flags))
-    elif case:
-        return pattern in self
+@contains.to_sql
+def _(self: sql.ColumnElement, substr: sql.ColumnElement, case: sql.ColumnElement | None = None) -> sql.ColumnElement:
+    # Replace all occurrences of `%`, `_`, and `\` with escaped versions
+    escaped_substr = sql.func.regexp_replace(substr, r'(%|_|\\)', r'\\\1', 'g')
+    if case is None:
+        # Default `case` is True, so we do a case-sensitive comparison
+        return self.like(sql.func.concat('%', escaped_substr, '%'))
     else:
-        return pattern.lower() in self.lower()
+        # Toggle case-sensitivity based on the value of `case`
+        return sql.case(
+            (case, self.like(sql.func.concat('%', escaped_substr, '%'))),
+            else_=sql.func.lower(self).like(sql.func.concat('%', sql.func.lower(escaped_substr), '%')),
+        )
+
+
+@pxt.udf(is_method=True)
+def contains_re(self: str, pattern: str, flags: int = 0) -> bool:
+    """
+    Test if string contains a regular expression pattern.
+
+    Args:
+        pattern: regular expression pattern
+        flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
+    """
+    return bool(re.search(pattern, self, flags))
 
 
 @pxt.udf(is_method=True)
@@ -84,22 +113,27 @@ def count(self: str, pattern: str, flags: int = 0) -> int:
         pattern: string literal or regular expression
         flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
     """
-    import re
-
     return builtins.len(re.findall(pattern, self, flags))
 
 
 @pxt.udf(is_method=True)
-def endswith(self: str, pattern: str) -> bool:
+def endswith(self: str, substr: str) -> bool:
     """
     Return `True` if the string ends with the specified suffix, otherwise return `False`.
 
     Equivalent to [`str.endswith()`](https://docs.python.org/3/library/stdtypes.html#str.endswith).
 
     Args:
-        pattern: string literal
+        substr: string literal
     """
-    return self.endswith(pattern)
+    return self.endswith(substr)
+
+
+@endswith.to_sql
+def _(self: sql.ColumnElement, substr: sql.ColumnElement) -> sql.ColumnElement:
+    # Replace all occurrences of `%`, `_`, and `\` with escaped versions
+    escaped_substr = sql.func.regexp_replace(substr, r'(%|_|\\)', r'\\\1', 'g')
+    return self.like(sql.func.concat('%', escaped_substr))
 
 
 @pxt.udf(is_method=True)
@@ -113,13 +147,11 @@ def fill(self: str, width: int, **kwargs: Any) -> str:
         width: Maximum line width.
         kwargs: Additional keyword arguments to pass to `textwrap.fill()`.
     """
-    import textwrap
-
     return textwrap.fill(self, width, **kwargs)
 
 
 @pxt.udf(is_method=True)
-def find(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def find(self: str, substr: str, start: int = 0, end: int | None = None) -> int:
     """
     Return the lowest index in string where `substr` is found within the slice `s[start:end]`.
 
@@ -133,6 +165,20 @@ def find(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] =
     return self.find(substr, start, end)
 
 
+@find.to_sql
+def _(
+    self: sql.ColumnElement, substr: sql.ColumnElement, start: sql.ColumnElement, end: sql.ColumnElement | None = None
+) -> sql.ColumnElement:
+    sl = pxt.functions.string.slice._to_sql(self, start, end)
+    if sl is None:
+        return None
+
+    strpos = sql.func.strpos(sl, substr)
+    return sql.case(
+        (strpos == 0, -1), (start >= 0, strpos + start - 1), else_=strpos + sql.func.char_length(self) + start - 1
+    )
+
+
 @pxt.udf(is_method=True)
 def findall(self: str, pattern: str, flags: int = 0) -> list:
     """
@@ -144,8 +190,6 @@ def findall(self: str, pattern: str, flags: int = 0) -> list:
         pattern: regular expression pattern
         flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
     """
-    import re
-
     return re.findall(pattern, self, flags)
 
 
@@ -171,8 +215,6 @@ def fullmatch(self: str, pattern: str, case: bool = True, flags: int = 0) -> boo
         case: if False, ignore case
         flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
     """
-    import re
-
     if not case:
         flags |= re.IGNORECASE
     _ = bool(re.fullmatch(pattern, self, flags))
@@ -180,7 +222,7 @@ def fullmatch(self: str, pattern: str, case: bool = True, flags: int = 0) -> boo
 
 
 @pxt.udf(is_method=True)
-def index(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def index(self: str, substr: str, start: int = 0, end: int | None = None) -> int:
     """
     Return the lowest index in string where `substr` is found within the slice `[start:end]`.
     Raises ValueError if `substr` is not found.
@@ -330,6 +372,11 @@ def len(self: str) -> int:
     return builtins.len(self)
 
 
+@len.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.char_length(self)
+
+
 @pxt.udf(is_method=True)
 def ljust(self: str, width: int, fillchar: str = ' ') -> str:
     """
@@ -355,8 +402,13 @@ def lower(self: str) -> str:
     return self.lower()
 
 
+@lower.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.lower(self)
+
+
 @pxt.udf(is_method=True)
-def lstrip(self: str, chars: Optional[str] = None) -> str:
+def lstrip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of the string with leading characters removed. The `chars` argument is a string specifying the set of
     characters to be removed. If omitted or `None`, whitespace characters are removed.
@@ -369,6 +421,11 @@ def lstrip(self: str, chars: Optional[str] = None) -> str:
     return self.lstrip(chars)
 
 
+@lstrip.to_sql
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
+    return sql.func.ltrim(self, chars if chars is not None else whitespace)
+
+
 @pxt.udf(is_method=True)
 def match(self: str, pattern: str, case: bool = True, flags: int = 0) -> bool:
     """
@@ -379,8 +436,6 @@ def match(self: str, pattern: str, case: bool = True, flags: int = 0) -> bool:
         case: if False, ignore case
         flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
     """
-    import re
-
     if not case:
         flags |= re.IGNORECASE
     return bool(re.match(pattern, self, flags))
@@ -440,9 +495,12 @@ def removeprefix(self: str, prefix: str) -> str:
     """
     Remove prefix. If the prefix is not present, returns string.
     """
-    if self.startswith(prefix):
-        return self[builtins.len(prefix) :]
-    return self
+    return self.removeprefix(prefix)
+
+
+@removeprefix.to_sql
+def _(self: sql.ColumnElement, prefix: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.case((startswith._to_sql(self, prefix), sql.func.right(self, -sql.func.char_length(prefix))), else_=self)
 
 
 @pxt.udf(is_method=True)
@@ -450,9 +508,12 @@ def removesuffix(self: str, suffix: str) -> str:
     """
     Remove suffix. If the suffix is not present, returns string.
     """
-    if self.endswith(suffix):
-        return self[: -builtins.len(suffix)]
-    return self
+    return self.removesuffix(suffix)
+
+
+@removesuffix.to_sql
+def _(self: sql.ColumnElement, suffix: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.case((endswith._to_sql(self, suffix), sql.func.left(self, -sql.func.char_length(suffix))), else_=self)
 
 
 @pxt.udf(is_method=True)
@@ -463,36 +524,69 @@ def repeat(self: str, n: int) -> str:
     return self * n
 
 
+@repeat.to_sql
+def _(self: sql.ColumnElement, n: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.repeat(self, n.cast(sql.types.INT))
+
+
 @pxt.udf(is_method=True)
-def replace(
-    self: str, pattern: str, repl: str, n: int = -1, case: bool = True, flags: int = 0, regex: bool = False
-) -> str:
+def replace(self: str, substr: str, repl: str, n: int | None = None) -> str:
     """
-    Replace occurrences of `pattern` with `repl`.
+    Replace occurrences of `substr` with `repl`.
 
-    Equivalent to [`str.replace()`](https://docs.python.org/3/library/stdtypes.html#str.replace) or
-    [`re.sub()`](https://docs.python.org/3/library/re.html#re.sub), depending on the value of regex.
+    Equivalent to [`str.replace()`](https://docs.python.org/3/library/stdtypes.html#str.replace).
 
     Args:
-        pattern: string literal or regular expression
+        substr: string literal
         repl: replacement string
-        n: number of replacements to make (-1 for all)
-        case: if False, ignore case
+        n: number of replacements to make (if `None`, replace all occurrences)
+    """
+    return self.replace(substr, repl, n or -1)
+
+
+@replace.to_sql
+def _(
+    self: sql.ColumnElement, substr: sql.ColumnElement, repl: sql.ColumnElement, n: sql.ColumnElement | None = None
+) -> sql.ColumnElement:
+    if n is not None:
+        return None  # SQL does not support bounding the number of replacements
+
+    return sql.func.replace(self, substr, repl)
+
+
+@pxt.udf(is_method=True)
+def replace_re(self: str, pattern: str, repl: str, n: int | None = None, flags: int = 0) -> str:
+    """
+    Replace occurrences of a regular expression pattern with `repl`.
+
+    Equivalent to [`re.sub()`](https://docs.python.org/3/library/re.html#re.sub).
+
+    Args:
+        pattern: regular expression pattern
+        repl: replacement string
+        n: number of replacements to make (if `None`, replace all occurrences)
         flags: [flags](https://docs.python.org/3/library/re.html#flags) for the `re` module
-        regex: if True, treat pattern as a regular expression
     """
-    if regex:
-        import re
+    return re.sub(pattern, repl, self, count=(n or 0), flags=flags)
+
+
+@pxt.udf(is_method=True)
+def reverse(self: str) -> str:
+    """
+    Return a reversed copy of the string.
+
+    Equivalent to `str[::-1]`.
+    """
+    return self[::-1]
 
-        if not case:
-            flags |= re.IGNORECASE
-        return re.sub(pattern, repl, self, count=(0 if n == -1 else n), flags=flags)
-    else:
-        return self.replace(pattern, repl, n)
+
+@reverse.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.reverse(self)
 
 
 @pxt.udf(is_method=True)
-def rfind(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def rfind(self: str, substr: str, start: int | None = 0, end: int | None = None) -> int:
     """
     Return the highest index where `substr` is found, such that `substr` is contained within `[start:end]`.
 
@@ -507,7 +601,7 @@ def rfind(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] =
 
 
 @pxt.udf(is_method=True)
-def rindex(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def rindex(self: str, substr: str, start: int | None = 0, end: int | None = None) -> int:
     """
     Return the highest index where `substr` is found, such that `substr` is contained within `[start:end]`.
     Raises ValueError if `substr` is not found.
@@ -544,7 +638,7 @@ def rpartition(self: str, sep: str = ' ') -> list:
 
 
 @pxt.udf(is_method=True)
-def rstrip(self: str, chars: Optional[str] = None) -> str:
+def rstrip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of string with trailing characters removed.
 
@@ -556,8 +650,13 @@ def rstrip(self: str, chars: Optional[str] = None) -> str:
     return self.rstrip(chars)
 
 
+@rstrip.to_sql
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
+    return sql.func.rtrim(self, chars if chars is not None else whitespace)
+
+
 @pxt.udf(is_method=True)
-def slice(self: str, start: Optional[int] = None, stop: Optional[int] = None, step: Optional[int] = None) -> str:
+def slice(self: str, start: int | None = None, stop: int | None = None, step: int | None = None) -> str:
     """
     Return a slice.
 
@@ -569,10 +668,43 @@ def slice(self: str, start: Optional[int] = None, stop: Optional[int] = None, st
     return self[start:stop:step]
 
 
-@pxt.udf(is_method=True)
-def slice_replace(
-    self: str, start: Optional[int] = None, stop: Optional[int] = None, repl: Optional[str] = None
-) -> str:
+@slice.to_sql
+def _(
+    self: sql.ColumnElement,
+    start: sql.ColumnElement | None = None,
+    stop: sql.ColumnElement | None = None,
+    step: sql.ColumnElement | None = None,
+) -> sql.ColumnElement:
+    if step is not None:
+        return None
+
+    if start is not None:
+        start = start.cast(sql.types.INT)  # Postgres won't accept a BIGINT
+        start = sql.case(
+            (start >= 0, start + 1),  # SQL is 1-based, Python is 0-based
+            else_=sql.func.char_length(self) + start + 1,  # negative index
+        )
+        start = sql.func.greatest(start, 1)
+
+    if stop is not None:
+        stop = stop.cast(sql.types.INT)  # Postgres won't accept a BIGINT
+        stop = sql.case(
+            (stop >= 0, stop + 1),  # SQL is 1-based, Python is 0-based
+            else_=sql.func.char_length(self) + stop + 1,  # negative index
+        )
+        stop = sql.func.greatest(stop, 0)
+
+    if start is None:
+        if stop is None:
+            return self
+        return sql.func.substr(self, 1, stop)
+    if stop is None:
+        return sql.func.substr(self, start)
+    return sql.func.substr(self, start, sql.func.greatest(stop - start, 0))
+
+
+@pxt.udf(is_method=True)
+def slice_replace(self: str, start: int | None = None, stop: int | None = None, repl: str | None = None) -> str:
     """
     Replace a positional slice with another value.
 
@@ -585,20 +717,27 @@ def slice_replace(
 
 
 @pxt.udf(is_method=True)
-def startswith(self: str, pattern: str) -> int:
+def startswith(self: str, substr: str) -> int:
     """
-    Return `True` if string starts with `pattern`, otherwise return `False`.
+    Return `True` if string starts with `substr`, otherwise return `False`.
 
     Equivalent to [`str.startswith()`](https://docs.python.org/3/library/stdtypes.html#str.startswith).
 
     Args:
-        pattern: string literal
+        substr: string literal
     """
-    return self.startswith(pattern)
+    return self.startswith(substr)
+
+
+@startswith.to_sql
+def _(self: sql.ColumnElement, substr: sql.ColumnElement) -> sql.ColumnElement:
+    # Replace all occurrences of `%`, `_`, and `\` with escaped versions
+    escaped_substr = sql.func.regexp_replace(substr, r'(%|_|\\)', r'\\\1', 'g')
+    return self.like(sql.func.concat(escaped_substr, '%'))
 
 
 @pxt.udf(is_method=True)
-def strip(self: str, chars: Optional[str] = None) -> str:
+def strip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of string with leading and trailing characters removed.
 
@@ -610,6 +749,11 @@ def strip(self: str, chars: Optional[str] = None) -> str:
     return self.strip(chars)
 
 
+@strip.to_sql
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
+    return sql.func.trim(self, chars if chars is not None else whitespace)
+
+
 @pxt.udf(is_method=True)
 def swapcase(self: str) -> str:
     """
@@ -641,6 +785,11 @@ def upper(self: str) -> str:
     return self.upper()
 
 
+@upper.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.upper(self)
+
+
 @pxt.udf(is_method=True)
 def wrap(self: str, width: int, **kwargs: Any) -> list[str]:
     """
@@ -653,8 +802,6 @@ def wrap(self: str, width: int, **kwargs: Any) -> list[str]:
         width: Maximum line width.
         kwargs: Additional keyword arguments to pass to `textwrap.fill()`.
     """
-    import textwrap
-
     return textwrap.wrap(self, width, **kwargs)
 
 
@@ -671,6 +818,29 @@ def zfill(self: str, width: int) -> str:
     return self.zfill(width)
 
 
+def string_splitter(text: Any, separators: str) -> tuple[type[pxt.iterators.ComponentIterator], dict[str, Any]]:
+    """Iterator over chunks of a string. The string is chunked according to the specified `separators`.
+
+    The iterator yields a `text` field containing the text of the chunk.
+    Chunked text will be cleaned with `ftfy.fix_text` to fix up common problems with unicode sequences.
+
+    Args:
+        separators: separators to use to chunk the document. Currently the only supported option is `'sentence'`.
+
+    Examples:
+        This example assumes an existing table `tbl` with a column `text` of type `pxt.String`.
+
+        Create a view that splits all strings on sentence boundaries:
+
+        >>> pxt.create_view(
+        ...     'sentence_chunks',
+        ...     tbl,
+        ...     iterator=string_splitter(tbl.text, separators='sentence')
+        ... )
+    """
+    return pxt.iterators.string.StringSplitter._create(text=text, separators=separators)
+
+
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/timestamp.py

@@ -1,5 +1,5 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `TimestampType`.
+Pixeltable UDFs for `TimestampType`.
 
 Usage example:
 ```python
@@ -11,7 +11,6 @@ t.select(t.timestamp_col.year, t.timestamp_col.weekday()).collect()
 """
 
 from datetime import datetime
-from typing import Optional
 
 import sqlalchemy as sql
 
@@ -134,7 +133,8 @@ def astimezone(self: datetime, tz: str) -> datetime:
     Convert the datetime to the given time zone.
 
     Args:
-        tz: The time zone to convert to. Must be a valid time zone name from the IANA Time Zone Database.
+        tz: The time zone to convert to. Must be a valid time zone name from the
+            [IANA Time Zone Database](https://en.wikipedia.org/wiki/List_of_tz_database_time_zones).
     """
     from zoneinfo import ZoneInfo
 
@@ -237,12 +237,12 @@ def _(
     microsecond: sql.ColumnElement = _SQL_ZERO,
 ) -> sql.ColumnElement:
     return sql.func.make_timestamptz(
-        sql.cast(year, sql.Integer),
-        sql.cast(month, sql.Integer),
-        sql.cast(day, sql.Integer),
-        sql.cast(hour, sql.Integer),
-        sql.cast(minute, sql.Integer),
-        sql.cast(second + microsecond / 1000000.0, sql.Float),
+        year.cast(sql.Integer),
+        month.cast(sql.Integer),
+        day.cast(sql.Integer),
+        hour.cast(sql.Integer),
+        minute.cast(sql.Integer),
+        (second + microsecond / 1000000.0).cast(sql.Float),
     )
 
 
@@ -271,13 +271,13 @@ def _(
 @pxt.udf(is_method=True)
 def replace(
     self: datetime,
-    year: Optional[int] = None,
-    month: Optional[int] = None,
-    day: Optional[int] = None,
-    hour: Optional[int] = None,
-    minute: Optional[int] = None,
-    second: Optional[int] = None,
-    microsecond: Optional[int] = None,
+    year: int | None = None,
+    month: int | None = None,
+    day: int | None = None,
+    hour: int | None = None,
+    minute: int | None = None,
+    second: int | None = None,
+    microsecond: int | None = None,
 ) -> datetime:
     """
     Return a datetime with the same attributes, except for those attributes given new values by whichever keyword