[REVIEW][FEA] Port nvtx.pyx to use non-legacy libcudf APIs

python/cudf/cudf/_libxx/cpp/nvtx.pxd

@@ -2,8 +2,10 @@
 
 from libc.stdint cimport uint32_t
 
+ctypedef uint32_t underlying_type_t_color
+
 cdef extern from "cudf/utilities/nvtx_utils.hpp" namespace "cudf::nvtx" nogil:
-    ctypedef enum _color 'color':
+    ctypedef enum color_type 'color':
         GREEN 'cudf::nvtx::color::GREEN'
         BLUE 'cudf::nvtx::color::BLUE'
         YELLOW 'cudf::nvtx::color::YELLOW'
@@ -14,14 +16,14 @@ cdef extern from "cudf/utilities/nvtx_utils.hpp" namespace "cudf::nvtx" nogil:
         DARK_GREEN 'cudf::nvtx::color::DARK_GREEN'
         ORANGE 'cudf::nvtx::color::ORANGE'
 
-    _color JOIN_COLOR 'cudf::nvtx::JOIN_COLOR'
-    _color GROUP_COLOR 'cudf::nvtx::GROUP_COLOR'
-    _color BINARY_OP_COLOR 'cudf::nvtx::BINARY_OP_COLOR'
-    _color PARTITION_COLOR 'cudf::nvtx::PARTITION_COLOR'
-    _color READ_CSV_COLOR 'cudf::nvtx::READ_CSV_COLOR'
+    color_type JOIN_COLOR 'cudf::nvtx::JOIN_COLOR'
+    color_type GROUP_COLOR 'cudf::nvtx::GROUP_COLOR'
+    color_type BINARY_OP_COLOR 'cudf::nvtx::BINARY_OP_COLOR'
+    color_type PARTITION_COLOR 'cudf::nvtx::PARTITION_COLOR'
+    color_type READ_CSV_COLOR 'cudf::nvtx::READ_CSV_COLOR'
 
     cdef void range_push(
-        const char* const name, _color color
+        const char* const name, color_type color
     ) except +
 
     cdef void range_push_hex(

python/cudf/cudf/_libxx/nvtx.pyx

@@ -1,15 +1,28 @@
 # Copyright (c) 2020, NVIDIA CORPORATION.
 
+from enum import IntEnum
 from libc.stdint cimport uint32_t
 
 from cudf._libxx.cpp.nvtx cimport (
     range_push as cpp_range_push,
     range_push_hex as cpp_range_push_hex,
     range_pop as cpp_range_pop,
 
-    _color as nvtx_color,
+    color_type,
+    underlying_type_t_color
+
 )
 
+class Color(IntEnum):
+    GREEN = <underlying_type_t_color> color_type.GREEN
+    BLUE = <underlying_type_t_color> color_type.BLUE
+    YELLOW = <underlying_type_t_color> color_type.YELLOW
+    PURPLE = <underlying_type_t_color> color_type.PURPLE
+    CYAN = <underlying_type_t_color> color_type.CYAN
+    RED = <underlying_type_t_color> color_type.RED
+    WHITE = <underlying_type_t_color> color_type.WHITE
+    DARK_GREEN = <underlying_type_t_color> color_type.DARK_GREEN
+    ORANGE = <underlying_type_t_color> color_type.ORANGE
 
 def range_push(object name, object color='green'):
     """
@@ -25,38 +38,14 @@ def range_push(object name, object color='green'):
     """
     cdef const char* _name = name
     cdef uint32_t _color = color
+    cdef color_type _Color = Color[color]
     try:
         _color = int(_color, 16)
         with nogil:
             cpp_range_push_hex(_name, _color)
     except ValueError:
-        if color == 'green':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.GREEN)
-        elif color == 'blue':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.BLUE)
-        elif color == 'yellow':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.YELLOW)
-        elif color == 'purple':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.PURPLE)
-        elif color == 'cyan':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.CYAN)
-        elif color == 'red':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.RED)
-        elif color == 'white':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.WHITE)
-        elif color == 'darkgreen':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.DARK_GREEN)
-        elif color == 'orange':
-            with nogil:
-                cpp_range_push(_name, nvtx_color.ORANGE)
+        with nogil:
+            cpp_range_push(_name, _Color)
 
 
 def range_pop():