Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[stdlib] Remove StringRef from python binding functions #3775

Open
wants to merge 5 commits into
base: nightly
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions stdlib/src/collections/string.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
These are Mojo built-ins, so you don't need to import them.
"""

from builtin.builtin_list import _lit_mut_cast
from collections import KeyElement, List, Optional
from collections._index_normalization import normalize_index
from sys import bitwidthof, llvm_intrinsic
Expand Down Expand Up @@ -1556,7 +1557,9 @@ struct String(
)

@always_inline
fn as_string_slice(ref self) -> StringSlice[__origin_of(self)]:
fn as_string_slice(
ref self,
) -> StringSlice[_lit_mut_cast[__origin_of(self), False].result]:
Copy link
Collaborator

Choose a reason for hiding this comment

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

Question Why force the return value to be an immutable slice?

Copy link
Contributor Author

@martinvuyk martinvuyk Nov 22, 2024

Choose a reason for hiding this comment

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

Because otherwise we'd need to add the equivalent of Span.get_immutable() and do string_var.as_string_slice().get_immutable() which is way too verbose. And I figured that to keep the scope of this PR small it might be better to just force the API to return immutable slices. I'm fine with doing it either way, WDYT?

"""Returns a string slice of the data owned by this string.

Returns:
Expand All @@ -1565,7 +1568,9 @@ struct String(
# FIXME(MSTDL-160):
# Enforce UTF-8 encoding in String so this is actually
# guaranteed to be valid.
return StringSlice(unsafe_from_utf8=self.as_bytes())
return StringSlice[_lit_mut_cast[__origin_of(self), False].result](
ptr=self.unsafe_ptr(), length=self.byte_length()
)

@always_inline
fn byte_length(self) -> Int:
Expand Down
147 changes: 83 additions & 64 deletions stdlib/src/python/_cpython.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ from python._bindings import Typed_initproc, PyMojoObject, Pythonable

from memory import UnsafePointer

from utils import StringRef, StringSlice
from utils import StringSlice


# ===-----------------------------------------------------------------------===#
Expand Down Expand Up @@ -260,7 +260,7 @@ struct PythonVersion:
"""The patch version number."""

@implicit
fn __init__(out self, version: StringRef):
fn __init__[O: ImmutableOrigin](out self, version: StringSlice[O]):
"""Initialize a PythonVersion object from a version string.

Args:
Expand Down Expand Up @@ -289,8 +289,16 @@ struct PythonVersion:
self = PythonVersion(components[0], components[1], components[2])


fn _py_get_version(lib: DLHandle) -> StringRef:
return StringRef(ptr=lib.call["Py_GetVersion", UnsafePointer[c_char]]())
fn _py_get_version(lib: DLHandle) -> StringSlice[ImmutableAnyOrigin]:
var ptr = lib.call["Py_GetVersion", UnsafePointer[c_char]]()
var length = 0
for i in range(50):
if ptr[i] == 0:
length = i
break
return StringSlice[ImmutableAnyOrigin](
ptr=ptr.bitcast[Byte](), length=length
)
martinvuyk marked this conversation as resolved.
Show resolved Hide resolved


fn _py_finalize(lib: DLHandle):
Expand Down Expand Up @@ -743,7 +751,7 @@ struct CPython:
"""The version of the Python runtime."""
var total_ref_count: UnsafePointer[Int]
"""The total reference count of all Python objects."""
var init_error: StringRef
var init_error: StringSlice[ImmutableAnyOrigin]
"""An error message if initialization failed."""

# ===-------------------------------------------------------------------===#
Expand Down Expand Up @@ -773,9 +781,17 @@ struct CPython:

# TODO(MOCO-772) Allow raises to propagate through function pointers
# and make this initialization a raising function.
self.init_error = external_call[
var ptr = external_call[
"KGEN_CompilerRT_Python_SetPythonPath", UnsafePointer[c_char]
]()
var length = 0
for i in range(1_000):
if ptr[i] == 0:
length = i
break
self.init_error = StringSlice[ImmutableAnyOrigin](
ptr=ptr.bitcast[Byte](), length=length
)
martinvuyk marked this conversation as resolved.
Show resolved Hide resolved

var python_lib = getenv("MOJO_PYTHON_LIBRARY")

Expand All @@ -789,7 +805,9 @@ struct CPython:
self.logging_enabled = logging_enabled
if not self.init_error:
if not self.lib.check_symbol("Py_Initialize"):
self.init_error = "compatible Python library not found"
self.init_error = __type_of(self.init_error)(
"compatible Python library not found"
)
self.lib.call["Py_Initialize"]()
self.version = PythonVersion(_py_get_version(self.lib))
else:
Expand Down Expand Up @@ -1062,15 +1080,16 @@ struct CPython:
# Python Module operations
# ===-------------------------------------------------------------------===#

fn PyImport_ImportModule(
inout self,
name: StringRef,
) -> PyObjectPtr:
fn PyImport_ImportModule[
O: ImmutableOrigin
](inout self, name: StringSlice[O]) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/import.html#c.PyImport_ImportModule).
"""

var r = self.lib.call["PyImport_ImportModule", PyObjectPtr](name.data)
var r = self.lib.call["PyImport_ImportModule", PyObjectPtr](
name.unsafe_ptr().bitcast[c_char]()
)

self.log(
r._get_ptr_as_int(),
Expand All @@ -1083,7 +1102,9 @@ struct CPython:
self._inc_total_rc()
return r

fn PyImport_AddModule(inout self, name: StringRef) -> PyObjectPtr:
fn PyImport_AddModule[
O: ImmutableOrigin
](inout self, name: StringSlice[O]) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/import.html#c.PyImport_AddModule).
"""
Expand Down Expand Up @@ -1190,7 +1211,9 @@ struct CPython:
# Python Evaluation
# ===-------------------------------------------------------------------===#

fn PyRun_SimpleString(inout self, strref: StringRef) -> Bool:
fn PyRun_SimpleString[
O: ImmutableOrigin
](inout self, strref: StringSlice[O]) -> Bool:
"""Executes the given Python code.

Args:
Expand All @@ -1205,12 +1228,17 @@ struct CPython:
https://docs.python.org/3/c-api/veryhigh.html#c.PyRun_SimpleString).
"""
return (
self.lib.call["PyRun_SimpleString", c_int](strref.unsafe_ptr()) == 0
self.lib.call["PyRun_SimpleString", c_int](
strref.unsafe_ptr().bitcast[c_char]()
)
== 0
)

fn PyRun_String(
fn PyRun_String[
O: ImmutableOrigin
](
inout self,
strref: StringRef,
strref: StringSlice[O],
globals: PyObjectPtr,
locals: PyObjectPtr,
run_mode: Int,
Expand All @@ -1219,7 +1247,10 @@ struct CPython:
https://docs.python.org/3/c-api/veryhigh.html#c.PyRun_String).
"""
var result = self.lib.call["PyRun_String", PyObjectPtr](
strref.unsafe_ptr(), Int32(run_mode), globals, locals
strref.unsafe_ptr().bitcast[c_char](),
Int32(run_mode),
globals,
locals,
)

self.log(
Expand Down Expand Up @@ -1256,18 +1287,22 @@ struct CPython:
"""
return self.lib.call["PyEval_GetBuiltins", PyObjectPtr]()

fn Py_CompileString(
fn Py_CompileString[
O1: ImmutableOrigin, O2: ImmutableOrigin
](
inout self,
strref: StringRef,
filename: StringRef,
strref: StringSlice[O1],
filename: StringSlice[O2],
compile_mode: Int,
) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/veryhigh.html#c.Py_CompileString).
"""

var r = self.lib.call["Py_CompileString", PyObjectPtr](
strref.unsafe_ptr(), filename.unsafe_ptr(), Int32(compile_mode)
strref.unsafe_ptr().bitcast[c_char](),
filename.unsafe_ptr().bitcast[c_char](),
Int32(compile_mode),
)
self._inc_total_rc()
return r
Expand Down Expand Up @@ -1353,17 +1388,15 @@ struct CPython:

return r

fn PyObject_GetAttrString(
inout self,
obj: PyObjectPtr,
name: StringRef,
) -> PyObjectPtr:
fn PyObject_GetAttrString[
O: ImmutableOrigin
](inout self, obj: PyObjectPtr, name: StringSlice[O],) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/object.html#c.PyObject_GetAttrString).
"""

var r = self.lib.call["PyObject_GetAttrString", PyObjectPtr](
obj, name.data
obj, name.unsafe_ptr().bitcast[c_char]()
)

self.log(
Expand All @@ -1379,15 +1412,20 @@ struct CPython:
self._inc_total_rc()
return r

fn PyObject_SetAttrString(
inout self, obj: PyObjectPtr, name: StringRef, new_value: PyObjectPtr
fn PyObject_SetAttrString[
O: ImmutableOrigin
](
inout self,
obj: PyObjectPtr,
name: StringSlice[O],
new_value: PyObjectPtr,
) -> c_int:
"""[Reference](
https://docs.python.org/3/c-api/object.html#c.PyObject_SetAttrString).
"""

var r = self.lib.call["PyObject_SetAttrString", c_int](
obj, name.data, new_value
obj, name.unsafe_ptr().bitcast[c_char](), new_value
)

self.log(
Expand All @@ -1404,9 +1442,7 @@ struct CPython:
return r

fn PyObject_CallObject(
inout self,
callable_obj: PyObjectPtr,
args: PyObjectPtr,
inout self, callable_obj: PyObjectPtr, args: PyObjectPtr
) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/call.html#c.PyObject_CallObject).
Expand Down Expand Up @@ -1705,15 +1741,17 @@ struct CPython:
# Unicode Objects
# ===-------------------------------------------------------------------===#

fn PyUnicode_DecodeUTF8(inout self, strref: StringRef) -> PyObjectPtr:
fn PyUnicode_DecodeUTF8[
O: ImmutableOrigin
](inout self, strref: StringSlice[O]) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_DecodeUTF8).
"""

var r = self.lib.call["PyUnicode_DecodeUTF8", PyObjectPtr](
strref.unsafe_ptr().bitcast[Int8](),
strref.length,
"strict".unsafe_cstr_ptr(),
strref.unsafe_ptr().bitcast[c_char](),
strref.byte_length(),
"strict".unsafe_ptr().bitcast[c_char](),
)

self.log(
Expand All @@ -1727,27 +1765,6 @@ struct CPython:
self._inc_total_rc()
return r

fn PyUnicode_DecodeUTF8(inout self, strslice: StringSlice) -> PyObjectPtr:
"""[Reference](
https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_DecodeUTF8).
"""
var r = self.lib.call["PyUnicode_DecodeUTF8", PyObjectPtr](
strslice.unsafe_ptr().bitcast[Int8](),
strslice.byte_length(),
"strict".unsafe_cstr_ptr(),
)

self.log(
r._get_ptr_as_int(),
" NEWREF PyUnicode_DecodeUTF8, refcnt:",
self._Py_REFCNT(r),
", str:",
strslice,
)

self._inc_total_rc()
return r

fn PySlice_FromSlice(inout self, slice: Slice) -> PyObjectPtr:
# Convert Mojo Slice to Python slice parameters
# Note: Deliberately avoid using `span.indices()` here and instead pass
Expand Down Expand Up @@ -1775,16 +1792,18 @@ struct CPython:

return py_slice

fn PyUnicode_AsUTF8AndSize(inout self, py_object: PyObjectPtr) -> StringRef:
fn PyUnicode_AsUTF8AndSize(
inout self, py_object: PyObjectPtr
) -> StringSlice[MutableAnyOrigin]:
"""[Reference](
https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_AsUTF8AndSize).
"""

var s = StringRef()
s.data = self.lib.call[
var length = 0
var ptr = self.lib.call[
"PyUnicode_AsUTF8AndSize", UnsafePointer[c_char]
](py_object, UnsafePointer.address_of(s.length)).bitcast[UInt8]()
return s
](py_object, UnsafePointer.address_of(length)).bitcast[Byte]()
return StringSlice[MutableAnyOrigin](ptr=ptr, length=length)

# ===-------------------------------------------------------------------===#
# Python Error operations
Expand Down
Loading