From 5c107744dad05ca6a8b446c5b1f6a049a00ce0f6 Mon Sep 17 00:00:00 2001
From: jakirkham <jakirkham@gmail.com>
Date: Fri, 27 Sep 2024 18:16:04 -0700
Subject: [PATCH] Touch ups to `Array`

---
 python/kvikio/kvikio/_lib/arr.pxd |  3 +++
 python/kvikio/kvikio/_lib/arr.pyi |  2 +-
 python/kvikio/kvikio/_lib/arr.pyx | 18 ++++++++++++++----
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/python/kvikio/kvikio/_lib/arr.pxd b/python/kvikio/kvikio/_lib/arr.pxd
index a9d2b7e7f1..47bad21a3b 100644
--- a/python/kvikio/kvikio/_lib/arr.pxd
+++ b/python/kvikio/kvikio/_lib/arr.pxd
@@ -28,6 +28,9 @@ cdef class Array:
     cpdef Py_ssize_t _nbytes(self)
 
 
+cpdef Array asarray(obj)
+
+
 cdef pair[uintptr_t, size_t] parse_buffer_argument(
     buf, size, bint accept_host_buffer
 ) except *
diff --git a/python/kvikio/kvikio/_lib/arr.pyi b/python/kvikio/kvikio/_lib/arr.pyi
index 0d2a8201eb..d2ba479895 100644
--- a/python/kvikio/kvikio/_lib/arr.pyi
+++ b/python/kvikio/kvikio/_lib/arr.pyi
@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 from typing import Generic, Tuple, TypeVar
diff --git a/python/kvikio/kvikio/_lib/arr.pyx b/python/kvikio/kvikio/_lib/arr.pyx
index 793a414b17..1d1afbe8d8 100644
--- a/python/kvikio/kvikio/_lib/arr.pyx
+++ b/python/kvikio/kvikio/_lib/arr.pyx
@@ -6,11 +6,20 @@
 
 from cpython.array cimport array, newarrayobject
 from cpython.buffer cimport PyBuffer_IsContiguous
-from cpython.memoryview cimport PyMemoryView_FromObject, PyMemoryView_GET_BUFFER
+from cpython.memoryview cimport (
+    PyMemoryView_FromObject,
+    PyMemoryView_GET_BUFFER,
+)
 from cpython.object cimport PyObject
 from cpython.ref cimport Py_INCREF
 from cpython.tuple cimport PyTuple_New, PyTuple_SET_ITEM
-from cython cimport auto_pickle, boundscheck, initializedcheck, nonecheck, wraparound
+from cython cimport (
+    auto_pickle,
+    boundscheck,
+    initializedcheck,
+    nonecheck,
+    wraparound,
+)
 from libc.stdint cimport uintptr_t
 from libc.string cimport memcpy
 
@@ -289,9 +298,10 @@ cdef inline Py_ssize_t _nbytes(Py_ssize_t itemsize,
         nbytes *= shape_mv[i]
     return nbytes
 
-cpdef asarray(obj):
+
+cpdef Array asarray(obj):
     if isinstance(obj, Array):
-        return obj
+        return <Array>obj
     else:
         return Array(obj)