Define the PythonBytes type.

Lukasa · Lukasa · commit 05179f5f649e · 2021-03-24T19:16:05.000Z
This type encapsulates the Python byte string object. This requires a
new data type, mostly because there is no natural "bucket of bytes" type
in Swift without adding a Foundation dependency, which this module
currently does not have.

Instead, we can define a little wrapper that makes it as easy as
possible to convert from existing types to the Python bytes object, and
back again.
diff --git a/PythonKit/Python.swift b/PythonKit/Python.swift
@@ -1430,3 +1430,90 @@ extension PythonObject : ExpressibleByArrayLiteral, ExpressibleByDictionaryLiter
         self.init(Dictionary(elements, uniquingKeysWith: { lhs, _ in lhs }))
     }
 }
+
+public struct PythonBytes : PythonConvertible, ConvertibleFromPython, Hashable {
+    public private(set) var pythonObject: PythonObject
+
+    public init?(_ pythonObject: PythonObject) {
+        // We try to get the string/size pointers out. If it works, hooray, this is a bytes
+        // otherwise it isn't.
+        let pyObject = pythonObject.ownedPyObject
+        defer { Py_DecRef(pyObject) }
+
+        var length = 0
+        var buffer: UnsafeMutablePointer<CChar>? = nil
+
+        switch PyBytes_AsStringAndSize(pyObject, &buffer, &length) {
+        case 0:
+            self.pythonObject = pythonObject
+        default:
+            return nil
+        }
+    }
+
+    @inlinable
+    public init<Bytes: Sequence>(_ bytes: Bytes) where Bytes.Element == UInt8 {
+        let possibleSelf = bytes.withContiguousStorageIfAvailable { storagePtr in
+            PythonBytes.fromBytePointer(storagePtr)
+        }
+        if let actualSelf = possibleSelf {
+            self = actualSelf
+        } else {
+            let temporaryBuffer = Array(bytes)
+            self = temporaryBuffer.withUnsafeBufferPointer {
+                PythonBytes.fromBytePointer($0)
+            }
+        }
+    }
+
+    @inlinable
+    public init<Bytes: Sequence>(_ bytes: Bytes) where Bytes.Element == Int8 {
+        let possibleSelf = bytes.withContiguousStorageIfAvailable { storagePtr in
+            PythonBytes.fromBytePointer(storagePtr)
+        }
+        if let actualSelf = possibleSelf {
+            self = actualSelf
+        } else {
+            let temporaryBuffer = Array(bytes)
+            self = temporaryBuffer.withUnsafeBufferPointer {
+                PythonBytes.fromBytePointer($0)
+            }
+        }
+    }
+
+    private init(bytesObject: PythonObject) {
+        self.pythonObject = bytesObject
+    }
+
+    @usableFromInline
+    static func fromBytePointer(_ bytes: UnsafeBufferPointer<UInt8>) -> PythonBytes {
+        bytes.withMemoryRebound(to: Int8.self) { reboundPtr in
+            PythonBytes.fromBytePointer(reboundPtr)
+        }
+    }
+
+    @usableFromInline
+    static func fromBytePointer(_ bytes: UnsafeBufferPointer<Int8>) -> PythonBytes {
+        let v = PyBytes_FromStringAndSize(bytes.baseAddress, bytes.count)!
+        return PythonBytes(bytesObject: PythonObject(consuming: v))
+    }
+
+    public func withUnsafeBytes<ReturnValue>(
+        _ callback: (UnsafeRawBufferPointer) throws -> ReturnValue
+    ) rethrows -> ReturnValue {
+        let pyObject = self.pythonObject.ownedPyObject
+        defer { Py_DecRef(pyObject) }
+
+        var length = 0
+        var buffer: UnsafeMutablePointer<CChar>? = nil
+
+        switch PyBytes_AsStringAndSize(pyObject, &buffer, &length) {
+        case 0:
+            let buffer = UnsafeRawBufferPointer(start: buffer, count: length)
+            return try callback(buffer)
+        default:
+            try! throwPythonErrorIfPresent()
+            fatalError("No result or error getting interior buffer for bytes \(self)")
+        }
+    }
+}
diff --git a/PythonKit/PythonLibrary+Symbols.swift b/PythonKit/PythonLibrary+Symbols.swift
@@ -185,6 +185,20 @@ let PyString_FromStringAndSize: @convention(c) (
         name: "PyUnicode_DecodeUTF8",
         legacyName: "PyString_FromStringAndSize")
 
+let PyBytes_FromStringAndSize: @convention(c) (
+    PyCCharPointer?, Int) -> (PyObjectPointer?) =
+    PythonLibrary.loadSymbol(
+        name: "PyBytes_FromStringAndSize",
+        legacyName: "PyString_FromStringAndSize")
+
+let PyBytes_AsStringAndSize: @convention(c) (
+    PyObjectPointer,
+    UnsafeMutablePointer<UnsafeMutablePointer<CChar>?>?,
+    UnsafeMutablePointer<Int>?) -> CInt =
+    PythonLibrary.loadSymbol(
+        name: "PyBytes_AsStringAndSize",
+        legacyName: "PyString_AsStringAndSize")
+
 let _Py_ZeroStruct: PyObjectPointer =
     PythonLibrary.loadSymbol(name: "_Py_ZeroStruct")
 
diff --git a/Tests/PythonKitTests/PythonRuntimeTests.swift b/Tests/PythonKitTests/PythonRuntimeTests.swift
@@ -279,4 +279,41 @@ class PythonRuntimeTests: XCTestCase {
             _ = Bool.init(b)
         }
     }
+
+    func testPythonBytes() {
+        let bytes = PythonBytes([UInt8(1), UInt8(2), UInt8(3), UInt8(4)])
+        bytes.withUnsafeBytes {
+            XCTAssertEqual(Array($0), [1, 2, 3, 4])
+        }
+    }
+
+    func testPythonBytesInt8() {
+        let bytes = PythonBytes([Int8(1), Int8(2), Int8(3), Int8(4)])
+        bytes.withUnsafeBytes {
+            XCTAssertEqual(Array($0), [1, 2, 3, 4])
+        }
+    }
+
+    func testPythonBytesNonContiguousSequence() {
+        let bytes = PythonBytes(CollectionOfOne(UInt8(1)))
+        bytes.withUnsafeBytes {
+            XCTAssertEqual(Array($0), [1])
+        }
+    }
+
+    func testPythonBytesNonContiguousSequenceInt8() {
+        let bytes = PythonBytes(CollectionOfOne(Int8(1)))
+        bytes.withUnsafeBytes {
+            XCTAssertEqual(Array($0), [1])
+        }
+    }
+
+    func testBytesConversion() {
+        let bytes = PythonBytes(CollectionOfOne(UInt8(1)))
+        let otherBytes = PythonBytes(bytes.pythonObject)
+        otherBytes?.withUnsafeBytes {
+            XCTAssertEqual(Array($0), [1])
+        }
+        XCTAssertEqual(bytes, otherBytes)
+    }
 }
