Skip to content
62 changes: 50 additions & 12 deletions numba/cpython/unicode.py
Original file line number Diff line number Diff line change
Expand Up @@ -813,18 +813,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 +830,56 @@ 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 (start is None or isinstance(start, (types.Omitted,
Comment thread
dlee992 marked this conversation as resolved.
Resolved
types.Integer,
types.NoneType))):
raise TypingError('When specified, the arg "start" must be an Integer')
Comment thread
dlee992 marked this conversation as resolved.
Resolved

if not (end is None or isinstance(end, (types.Omitted,
types.Integer,
types.NoneType))):
raise TypingError('When specified, the arg "end" must be an Integer')
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!

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 = len(s)
sub_length = len(prefix)
if start is None:
start = 0
if end is None:
end = length

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

if sub_length == 0:
return True

s = s[start:end]

return _cmp_region(s, 0, prefix, 0, sub_length) == 0
Comment thread
dlee992 marked this conversation as resolved.
Resolved

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
113 changes: 100 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,102 @@ 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_start = 'When specified, the arg "start" must be an Integer'
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'
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