Use a local variable for conditional branching

cupy · May 22, 2018 · 6683e5f · 6683e5f
1 parent 3029e10
commit 6683e5f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/cupy/core/dlpack.pyx b/cupy/core/dlpack.pyx
@@ -197,34 +197,35 @@ cpdef ndarray fromDlpack(object dltensor):
     mem = DLPackMemory(dltensor)
 
     cdef DLDataType dtype = mem.dlm_tensor.dl_tensor.dtype
+    cdef int bits = dtype.bits
     if dtype.code == DLDataTypeCode.kDLUInt:
-        if dtype.bits == 8:
+        if bits == 8:
             cp_dtype = cupy.uint8
-        elif dtype.bits == 16:
+        elif bits == 16:
             cp_dtype = cupy.uint16
-        elif dtype.bits == 32:
+        elif bits == 32:
             cp_dtype = cupy.uint32
-        elif dtype.bits == 64:
+        elif bits == 64:
             cp_dtype = cupy.uint64
         else:
-            raise TypeError('uint{} is not supported.'.format(dtype.bits))
+            raise TypeError('uint{} is not supported.'.format(bits))
     elif dtype.code == DLDataTypeCode.kDLInt:
-        if dtype.bits == 8:
+        if bits == 8:
             cp_dtype = cupy.int8
-        elif dtype.bits == 16:
+        elif bits == 16:
             cp_dtype = cupy.int16
-        elif dtype.bits == 32:
+        elif bits == 32:
             cp_dtype = cupy.int32
-        elif dtype.bits == 64:
+        elif bits == 64:
             cp_dtype = cupy.int64
         else:
-            raise TypeError('int{} is not supported.'.format(dtype.bits))
+            raise TypeError('int{} is not supported.'.format(bits))
     elif dtype.code == DLDataTypeCode.kDLFloat:
-        if dtype.bits == 16:
+        if bits == 16:
             cp_dtype = cupy.float16
-        elif dtype.bits == 32:
+        elif bits == 32:
             cp_dtype = cupy.float32
-        elif dtype.bits == 64:
+        elif bits == 64:
             cp_dtype = cupy.float64
         else:
             raise TypeError('float{} is not supported.'.format(dtype.bits))

diff --git a/tests/cupy_tests/core_tests/test_dlpack.py b/tests/cupy_tests/core_tests/test_dlpack.py
@@ -27,6 +27,7 @@ def test_conversion(self):
         tensor = self.array.toDlpack()
         array = cupy.fromDlpack(tensor)
         testing.assert_array_equal(self.array, array)
+        testing.assert_array_equal(self.array.data.ptr, array.data.ptr)
 
 
 class TestDLTensorMemory(unittest.TestCase):