Skip to content
49 changes: 26 additions & 23 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 @@ -831,51 +832,53 @@ def _adjust_indices(length, start, end):


@overload_method(types.UnicodeType, 'startswith')
def unicode_startswith(s, substr, start=None, end=None):
if not (start is None or isinstance(start, (types.Omitted,
types.Integer,
types.NoneType))):
raise TypingError('When specified, the arg "start" must be an Integer')
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 (end is None or isinstance(end, (types.Omitted,
types.Integer,
types.NoneType))):
raise TypingError('When specified, the arg "end" must be an Integer')
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(substr, (types.Tuple, types.UniTuple)):
def startswith_tuple_impl(s, substr, start=None, end=None):
for item in substr:
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(substr, types.UnicodeCharSeq):
def startswith_char_seq_impl(s, substr, start=None, end=None):
return s.startswith(str(substr), start, end)
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(substr, types.UnicodeType):
def startswith_unicode_impl(s, substr, start=None, end=None):
length = len(s)
sub_length = len(substr)
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 < sub_length:
if end - start < prefix_length:
return False

if sub_length == 0:
if prefix_length == 0:
return True

s = s[start:end]
s_slice = s[start:end]

return _cmp_region(s, 0, substr, 0, sub_length) == 0
return _cmp_region(s_slice, 0, prefix, 0, prefix_length) == 0

return startswith_unicode_impl

Comment thread
dlee992 marked this conversation as resolved.
Resolved
Expand Down
50 changes: 32 additions & 18 deletions numba/tests/test_unicode.py
Original file line number Diff line number Diff line change
Expand Up @@ -571,8 +571,8 @@ def test_startswith_default(self):
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:
self.assertEqual(pyfunc(s, sub_str), cfunc(s, sub_str))
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
Expand All @@ -586,12 +586,10 @@ def test_startswith_with_start(self):
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 prefix in cpython_subs + default_subs + extra_subs:
for start in list(range(-20, 20)) + [None]:
self.assertEqual(pyfunc(s, sub_str, start),
cfunc(s, sub_str, start))
with self.assertRaises(TypingError):
cfunc('hello', 'he', 0.1)
self.assertEqual(pyfunc(s, prefix, start),
cfunc(s, prefix, start))

def test_startswith_with_start_end(self):
pyfunc = startswith_with_start_end_usecase
Expand All @@ -605,13 +603,29 @@ def test_startswith_with_start_end(self):
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 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, sub_str, start, end),
cfunc(s, sub_str, start, end))
with self.assertRaises(TypingError):
cfunc('hello', 'he', 0, 0.1)
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
Expand All @@ -626,9 +640,9 @@ def test_startswith_tuple(self):
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:
tuple_subs = (sub_str, 'lo')
self.assertEqual(pyfunc(s, tuple_subs),
cfunc(s, tuple_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
Expand All @@ -645,9 +659,9 @@ def test_startswith_tuple_args(self):
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]:
tuple_subs = (sub_str, 'lo')
self.assertEqual(pyfunc(s, tuple_subs, start, end),
cfunc(s, tuple_subs, start, end))
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