Skip to content
65 changes: 53 additions & 12 deletions numba/cpython/unicode.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import numpy as np
from llvmlite.ir import IntType, Constant

from numba.core.cgutils import is_nonelike
from numba.core.extending import (
models,
register_model,
Expand Down Expand Up @@ -813,18 +814,6 @@ def impl(data, sep):
return impl


@overload_method(types.UnicodeType, 'startswith')
def unicode_startswith(a, b):
if isinstance(b, types.UnicodeType):
def startswith_impl(a, b):
return _cmp_region(a, 0, b, 0, len(b)) == 0
return startswith_impl
if isinstance(b, types.UnicodeCharSeq):
def startswith_impl(a, b):
return a.startswith(str(b))
return startswith_impl


# https://github.com/python/cpython/blob/201c8f79450628241574fba940e08107178dc3a5/Objects/unicodeobject.c#L9342-L9354 # noqa: E501
@register_jitable
def _adjust_indices(length, start, end):
Expand All @@ -842,6 +831,58 @@ def _adjust_indices(length, start, end):
return start, end


@overload_method(types.UnicodeType, 'startswith')
def unicode_startswith(s, prefix, start=None, end=None):
if not is_nonelike(start) and not isinstance(start, types.Integer):
raise TypingError(
"When specified, the arg 'start' must be an Integer or nonelike")
Comment thread
dlee992 marked this conversation as resolved.
Resolved

if not is_nonelike(end) and not isinstance(end, types.Integer):
raise TypingError(
"When specified, the arg 'end' must be an Integer or nonelike")
Comment thread
dlee992 marked this conversation as resolved.
Resolved

if isinstance(prefix, (types.Tuple, types.UniTuple)):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tuple cannot be accepted here unless the contents of the Tuple are all types.StringLiteral. Consider:

@njit
def foo():
    print('abcde'.startswith(('a', 1)))

foo()

it may be easier to just not accept Tuple to force type inference to try with a non-literal types instead.

@dlee992 dlee992 Nov 4, 2022 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I prefer to accept a UniTuple of UnicodeType, since it's more consistent with cpython. So I want to add a check to ensure the prefix type as expected. BTW, in your case:

@njit
def foo():
    print('abcde'.startswith(('a', 1)))

or another case:

prefix = ('a', 1)
@njit
def foo(prefix):
    print('abcde'.startswith(prefix))

Both ('a', 1) will be inferred as Tuple(unicode_type, int64), than Tuple containing a StringLiteral.
And in the executable case:

@njit
def foo():
    print('abcde'.startswith(('a', 'b')))

('a', 'b') will be inferred as UniTuple(unicode_type x 2), so I just have to ensure prefix is a UniTuple with unicode_type dtype. If I'm wrong, please point out with an example. I haven't found a counterexample...

BTW, I found an interesting thing in cpython (which I think we don't have to implement this short-cut strategy):

>>> 'a'.startswith(('a', 5))
True

>>> 'a'.startswith((1, 'a'))
Traceback (most recent call last):
  File "/Users/ld/opt/anaconda3/envs/numbaenv/lib/python3.10/code.py", line 90, in runcode
    exec(code, self.locals)
  File "<input>", line 1, in <module>
TypeError: tuple for startswith must only contain str, not int

@dlee992 dlee992 Nov 4, 2022 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Revised, please review this issue again. @stuartarchibald

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By default, type inference will try using non-literal types first, and then if that fails, it will try again using literal types. Python source containing tuples of string literals will be inferred as types.Tuple(<types.StringLiteral>, <types.StringLiteral>, ...). e.g.

from numba import njit
import numpy as np

@njit
def foo():
    x = ('a', 'b')

foo()
foo.inspect_types()

You can see the "try non-literal then literal" types behaviour by doing something like:

from numba import njit, types
from numba.extending import overload
import numpy as np

def bar(tup):
    pass

@overload(bar)
def ol_bar(tup):
    print(f"bar:: tup type is {tup}")

@njit
def foo():
    x = ('a', 'b')
    bar(x)

foo()

(ignore the compilation error, just look at what gets printed to stdout).


RE: the 'a'.startswith(('a', 5)), there are implementation details like this in the CPython implementation. They don't need replicating unless there's some proven use case!

if isinstance(prefix, types.Tuple) or \
not isinstance(prefix.dtype, types.UnicodeType):
raise TypingError(
"The arg 'prefix' should a UniTuple of UnicodeType")
Comment thread
dlee992 marked this conversation as resolved.
Resolved

def startswith_tuple_impl(s, prefix, start=None, end=None):
for item in prefix:
if s.startswith(item, start, end):
return True
return False

return startswith_tuple_impl

if isinstance(prefix, types.UnicodeCharSeq):
Comment thread
dlee992 marked this conversation as resolved.
Resolved
def startswith_char_seq_impl(s, prefix, start=None, end=None):
return s.startswith(str(prefix), start, end)

return startswith_char_seq_impl

if isinstance(prefix, types.UnicodeType):
Comment thread
dlee992 marked this conversation as resolved.
Resolved
def startswith_unicode_impl(s, prefix, start=None, end=None):
length, prefix_length = len(s), len(prefix)
if start is None:
start = 0
if end is None:
end = length

start, end = _adjust_indices(length, start, end)
if end - start < prefix_length:
return False

if prefix_length == 0:
return True

s_slice = s[start:end]

return _cmp_region(s_slice, 0, prefix, 0, prefix_length) == 0

return startswith_unicode_impl

Comment thread
dlee992 marked this conversation as resolved.
Resolved

@overload_method(types.UnicodeType, 'endswith')
def unicode_endswith(s, substr, start=None, end=None):
if not (start is None or isinstance(start, (types.Omitted,
Expand Down
120 changes: 107 additions & 13 deletions numba/tests/test_unicode.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,14 @@ def expandtabs_with_tabsize_kwarg_usecase(s, tabsize):
return s.expandtabs(tabsize=tabsize)


def startswith_with_start_only_usecase(x, y, start):
return x.startswith(y, start)


def startswith_with_start_end_usecase(x, y, start, end):
return x.startswith(y, start, end)


def endswith_with_start_only_usecase(x, y, start):
return x.endswith(y, start)

Expand Down Expand Up @@ -551,23 +559,109 @@ def test_expandtabs_exception_noninteger_tabsize(self):
msg = '"tabsize" must be {}, not float'.format(accepted_types)
self.assertIn(msg, str(raises.exception))

def test_startswith(self, flags=no_pyobj_flags):
def test_startswith_default(self):
pyfunc = startswith_usecase
cfunc = njit(pyfunc)
for a in UNICODE_EXAMPLES:
for b in ['', 'x', a[:-2], a[3:], a, a + a]:
self.assertEqual(pyfunc(a, b),
cfunc(a, b),
'%s, %s' % (a, b))

def test_endswith(self, flags=no_pyobj_flags):
pyfunc = endswith_usecase
cpython_str = ['hello', 'helloworld', '']
cpython_subs = [
'he', 'hello', 'helloworld', 'ello',
'', 'lowo', 'lo', 'he', 'lo', 'o',
]
extra_subs = ['hellohellohello', ' ']
for s in cpython_str + UNICODE_EXAMPLES:
default_subs = ['', 'x', s[:-2], s[3:], s, s + s]
for prefix in cpython_subs + default_subs + extra_subs:
self.assertEqual(pyfunc(s, prefix), cfunc(s, prefix))

def test_startswith_with_start(self):
pyfunc = startswith_with_start_only_usecase
cfunc = njit(pyfunc)
for a in UNICODE_EXAMPLES:
for b in ['', 'x', a[:-2], a[3:], a, a + a]:
self.assertEqual(pyfunc(a, b),
cfunc(a, b),
'%s, %s' % (a, b))

cpython_str = ['hello', 'helloworld', '']
cpython_subs = [
'he', 'hello', 'helloworld', 'ello',
'', 'lowo', 'lo', 'he', 'lo', 'o',
]
extra_subs = ['hellohellohello', ' ']
for s in cpython_str + UNICODE_EXAMPLES:
default_subs = ['', 'x', s[:-2], s[3:], s, s + s]
for prefix in cpython_subs + default_subs + extra_subs:
for start in list(range(-20, 20)) + [None]:
self.assertEqual(pyfunc(s, prefix, start),
cfunc(s, prefix, start))

def test_startswith_with_start_end(self):
pyfunc = startswith_with_start_end_usecase
cfunc = njit(pyfunc)

cpython_str = ['hello', 'helloworld', '']
cpython_subs = [
'he', 'hello', 'helloworld', 'ello',
'', 'lowo', 'lo', 'he', 'lo', 'o',
]
extra_subs = ['hellohellohello', ' ']
for s in cpython_str + UNICODE_EXAMPLES:
default_subs = ['', 'x', s[:-2], s[3:], s, s + s]
for prefix in cpython_subs + default_subs + extra_subs:
for start in list(range(-20, 20)) + [None]:
for end in list(range(-20, 20)) + [None]:
self.assertEqual(pyfunc(s, prefix, start, end),
cfunc(s, prefix, start, end))

def test_startswith_exception_invalid_args(self):
msg_invalid_prefix = "The arg 'prefix' should a UniTuple of UnicodeType"
with self.assertRaisesRegex(TypingError, msg_invalid_prefix):
cfunc = njit(startswith_usecase)
cfunc("hello", (1, "he"))

msg_invalid_start = \
"When specified, the arg 'start' must be an Integer or nonelike"
with self.assertRaisesRegex(TypingError, msg_invalid_start):
cfunc = njit(startswith_with_start_only_usecase)
cfunc("hello", "he", "invalid start")

msg_invalid_end = \
"When specified, the arg 'end' must be an Integer or nonelike"
with self.assertRaisesRegex(TypingError, msg_invalid_end):
cfunc = njit(startswith_with_start_end_usecase)
cfunc("hello", "he", 0, "invalid end")

def test_startswith_tuple(self):
pyfunc = startswith_usecase
cfunc = njit(pyfunc)

cpython_str = ['hello', 'helloworld', '']
cpython_subs = [
'he', 'hello', 'helloworld', 'ello',
'', 'lowo', 'lo', 'he', 'lo', 'o',
]
extra_subs = ['hellohellohello', ' ']
for s in cpython_str + UNICODE_EXAMPLES:
default_subs = ['', 'x', s[:-2], s[3:], s, s + s]
for sub_str in cpython_subs + default_subs + extra_subs:
prefix = (sub_str, 'lo')
self.assertEqual(pyfunc(s, prefix),
cfunc(s, prefix))

def test_startswith_tuple_args(self):
pyfunc = startswith_with_start_end_usecase
cfunc = njit(pyfunc)

cpython_str = ['hello', 'helloworld', '']
cpython_subs = [
'he', 'hello', 'helloworld', 'ello',
'', 'lowo', 'lo', 'he', 'lo', 'o',
]
extra_subs = ['hellohellohello', ' ']
for s in cpython_str + UNICODE_EXAMPLES:
default_subs = ['', 'x', s[:-2], s[3:], s, s + s]
for sub_str in cpython_subs + default_subs + extra_subs:
for start in list(range(-20, 20)) + [None]:
for end in list(range(-20, 20)) + [None]:
prefix = (sub_str, 'lo')
self.assertEqual(pyfunc(s, prefix, start, end),
cfunc(s, prefix, start, end))

def test_endswith_default(self):
pyfunc = endswith_usecase
Expand Down