Enabling enum flags plugin changes (#2114)

rahulr-NI · web-flow · commit 37ea6ef997af · 2025-07-14T13:21:09.000-05:00
* Enabling Enum Flags plugin changes

* Updating the test cases

* Making import pytest to global

* Updating documentation changes

* Updating enum type

* Updating the test case to check the values

* Updating the test case based on code review

* Code review comments

* Including comment

* Updated test cases
diff --git a/build/templates/_visatype.py b/build/templates/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/build/templates/enums.py.mako b/build/templates/enums.py.mako
@@ -6,10 +6,13 @@ config = template_parameters['metadata'].config
 enums = config['enums']
 %>
 from enum import Enum
+% if any(enums[e].get('enum_class', 'Enum') == 'IntFlag' for e in enums):
+from enum import IntFlag
+% endif
 % for enum_name in sorted(helper.filter_codegen_enums(enums)):
 
 
-class ${enums[enum_name]['python_name']}(Enum):
+class ${enums[enum_name]['python_name']}(${enums[enum_name].get('enum_class', 'Enum')}):
 <%
     print_list = []
 %>\
diff --git a/generated/nidcpower/nidcpower/_visatype.py b/generated/nidcpower/nidcpower/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nidigital/nidigital/_visatype.py b/generated/nidigital/nidigital/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nidmm/nidmm/_visatype.py b/generated/nidmm/nidmm/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nifake/nifake/_grpc_stub_interpreter.py b/generated/nifake/nifake/_grpc_stub_interpreter.py
@@ -144,6 +144,12 @@ def fetch_waveform(self, number_of_samples):  # noqa: N802
     def fetch_waveform_into(self, number_of_samples):  # noqa: N802
         raise NotImplementedError('numpy-specific methods are not supported over gRPC')
 
+    def function_with_intflag_parameter(self, flag):  # noqa: N802
+        self._invoke(
+            self._client.FunctionWithIntflagParameter,
+            grpc_types.FunctionWithIntflagParameterRequest(vi=self._vi, flag=flag.value),
+        )
+
     def function_with_repeated_capability_type(self, site_list):  # noqa: N802
         raise NotImplementedError('function_with_repeated_capability_type is not supported over gRPC')
 
diff --git a/generated/nifake/nifake/_library.py b/generated/nifake/nifake/_library.py
@@ -36,6 +36,7 @@ def __init__(self, ctypes_library):
         self.niFake_EnumInputFunctionWithDefaults_cfunc = None
         self.niFake_ExportAttributeConfigurationBuffer_cfunc = None
         self.niFake_FetchWaveform_cfunc = None
+        self.niFake_FunctionWithIntflagParameter_cfunc = None
         self.niFake_FunctionWithRepeatedCapabilityType_cfunc = None
         self.niFake_GetABoolean_cfunc = None
         self.niFake_GetANumber_cfunc = None
@@ -187,6 +188,14 @@ def niFake_FetchWaveform(self, vi, number_of_samples, waveform_data, actual_numb
                 self.niFake_FetchWaveform_cfunc.restype = ViStatus  # noqa: F405
         return self.niFake_FetchWaveform_cfunc(vi, number_of_samples, waveform_data, actual_number_of_samples)
 
+    def niFake_FunctionWithIntflagParameter(self, vi, flag):  # noqa: N802
+        with self._func_lock:
+            if self.niFake_FunctionWithIntflagParameter_cfunc is None:
+                self.niFake_FunctionWithIntflagParameter_cfunc = self._get_library_function('niFake_FunctionWithIntflagParameter')
+                self.niFake_FunctionWithIntflagParameter_cfunc.argtypes = [ViSession, ViUInt64]  # noqa: F405
+                self.niFake_FunctionWithIntflagParameter_cfunc.restype = ViStatus  # noqa: F405
+        return self.niFake_FunctionWithIntflagParameter_cfunc(vi, flag)
+
     def niFake_FunctionWithRepeatedCapabilityType(self, vi, site_list):  # noqa: N802
         with self._func_lock:
             if self.niFake_FunctionWithRepeatedCapabilityType_cfunc is None:
diff --git a/generated/nifake/nifake/_library_interpreter.py b/generated/nifake/nifake/_library_interpreter.py
@@ -218,6 +218,13 @@ def fetch_waveform_into(self, waveform_data):  # noqa: N802
         errors.handle_error(self, error_code, ignore_warnings=False, is_error_handling=False)
         return
 
+    def function_with_intflag_parameter(self, flag):  # noqa: N802
+        vi_ctype = _visatype.ViSession(self._vi)  # case S110
+        flag_ctype = _visatype.ViUInt64(flag.value)  # case S130
+        error_code = self._library.niFake_FunctionWithIntflagParameter(vi_ctype, flag_ctype)
+        errors.handle_error(self, error_code, ignore_warnings=False, is_error_handling=False)
+        return
+
     def function_with_repeated_capability_type(self, site_list):  # noqa: N802
         vi_ctype = _visatype.ViSession(self._vi)  # case S110
         site_list_ctype = ctypes.create_string_buffer(site_list.encode(self._encoding))  # case C010
diff --git a/generated/nifake/nifake/_visatype.py b/generated/nifake/nifake/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nifake/nifake/enums.py b/generated/nifake/nifake/enums.py
@@ -2,6 +2,7 @@
 # This file was generated
 
 from enum import Enum
+from enum import IntFlag
 
 
 class Color(Enum):
@@ -53,6 +54,21 @@ class FloatEnum(Enum):
     '''
 
 
+class IntFlagEnum(IntFlag):
+    A = 1
+    r'''
+    Flag A option.
+    '''
+    B = 2
+    r'''
+    Flag B option.
+    '''
+    C = 9223372036854775808
+    r'''
+    Flag C option.
+    '''
+
+
 class MobileOSNames(Enum):
     ANDROID = 'Android'
     r'''
diff --git a/generated/nifake/nifake/session.py b/generated/nifake/nifake/session.py
@@ -945,6 +945,20 @@ def fetch_waveform_into(self, waveform_data):
             raise TypeError('waveform_data must be numpy.ndarray of dtype=float64, is ' + str(waveform_data.dtype))
         self._interpreter.fetch_waveform_into(waveform_data)
 
+    @ivi_synchronized
+    def function_with_intflag_parameter(self, flag):
+        r'''function_with_intflag_parameter
+
+        Calls a method that takes a flag parameter which can be OR'd from multiple enum values.
+
+        Args:
+            flag (enums.IntFlagEnum): A flag parameter that can be a combination (bitwise OR) of IntFlagEnum values.
+
+        '''
+        if type(flag) is not enums.IntFlagEnum:
+            raise TypeError('Parameter flag must be of type ' + str(enums.IntFlagEnum))
+        self._interpreter.function_with_intflag_parameter(flag)
+
     @ivi_synchronized
     def get_a_boolean(self):
         r'''get_a_boolean
diff --git a/generated/nifake/nifake/unit_tests/_mock_helper.py b/generated/nifake/nifake/unit_tests/_mock_helper.py
@@ -42,6 +42,8 @@ def __init__(self):
         self._defaults['FetchWaveform']['return'] = 0
         self._defaults['FetchWaveform']['waveformData'] = None
         self._defaults['FetchWaveform']['actualNumberOfSamples'] = None
+        self._defaults['FunctionWithIntflagParameter'] = {}
+        self._defaults['FunctionWithIntflagParameter']['return'] = 0
         self._defaults['FunctionWithRepeatedCapabilityType'] = {}
         self._defaults['FunctionWithRepeatedCapabilityType']['return'] = 0
         self._defaults['GetABoolean'] = {}
@@ -337,6 +339,11 @@ def niFake_FetchWaveform(self, vi, number_of_samples, waveform_data, actual_numb
             actual_number_of_samples.contents.value = self._defaults['FetchWaveform']['actualNumberOfSamples']
         return self._defaults['FetchWaveform']['return']
 
+    def niFake_FunctionWithIntflagParameter(self, vi, flag):  # noqa: N802
+        if self._defaults['FunctionWithIntflagParameter']['return'] != 0:
+            return self._defaults['FunctionWithIntflagParameter']['return']
+        return self._defaults['FunctionWithIntflagParameter']['return']
+
     def niFake_FunctionWithRepeatedCapabilityType(self, vi, site_list):  # noqa: N802
         if self._defaults['FunctionWithRepeatedCapabilityType']['return'] != 0:
             return self._defaults['FunctionWithRepeatedCapabilityType']['return']
@@ -1036,6 +1043,8 @@ def set_side_effects_and_return_values(self, mock_library):
         mock_library.niFake_ExportAttributeConfigurationBuffer.return_value = 0
         mock_library.niFake_FetchWaveform.side_effect = MockFunctionCallError("niFake_FetchWaveform")
         mock_library.niFake_FetchWaveform.return_value = 0
+        mock_library.niFake_FunctionWithIntflagParameter.side_effect = MockFunctionCallError("niFake_FunctionWithIntflagParameter")
+        mock_library.niFake_FunctionWithIntflagParameter.return_value = 0
         mock_library.niFake_FunctionWithRepeatedCapabilityType.side_effect = MockFunctionCallError("niFake_FunctionWithRepeatedCapabilityType")
         mock_library.niFake_FunctionWithRepeatedCapabilityType.return_value = 0
         mock_library.niFake_GetABoolean.side_effect = MockFunctionCallError("niFake_GetABoolean")
diff --git a/generated/nifake/nifake/unit_tests/test_session.py b/generated/nifake/nifake/unit_tests/test_session.py
@@ -4,6 +4,7 @@
 import nifake
 import nifake.errors
 import numpy
+import pytest
 
 from unittest.mock import MagicMock
 from unittest.mock import patch
@@ -837,10 +838,26 @@ def test_return_timedeltas(self):
             assert returned_timedeltas == expected_timedeltas
             self.patched_library_interpreter.return_list_of_durations_in_seconds.assert_called_once_with(len(time_values))
 
+    def test_with_valid_intflag_parameter(self):
+        flags = nifake.IntFlagEnum.C | nifake.IntFlagEnum.A
+        self.patched_library_interpreter.function_with_intflag_parameter.side_effect = None
+        self.patched_library_interpreter.function_with_intflag_parameter.return_value = None
+        with nifake.Session('dev1') as session:
+            session.function_with_intflag_parameter(flags)
+        self.patched_library_interpreter.function_with_intflag_parameter.assert_called_once_with(9223372036854775809)
+
+    def test_with_intflag_parameter_invalid(self):
+        invalid_flag = 5
+        with nifake.Session('dev1') as session:
+            try:
+                session.function_with_intflag_parameter(invalid_flag)
+                assert False
+            except TypeError:
+                pass
+
     def test_session_write_waveform_numpy_complex64_invalid_dtype(self):
         invalid_waveform_data = numpy.full(10, 1.0 + 1.0j, dtype=numpy.complex128)
         expected_error_message = "waveform_data_array must be numpy.ndarray of dtype=complex64, is complex128"
-        import pytest
         with nifake.Session('dev1') as session:
             with pytest.raises(TypeError) as exc_info:
                 session.write_waveform_numpy_complex64(invalid_waveform_data)
@@ -849,7 +866,6 @@ def test_session_write_waveform_numpy_complex64_invalid_dtype(self):
     def test_session_write_waveform_numpy_complex128_invalid_dtype(self):
         invalid_waveform_data = numpy.full(10, 1.0 + 1.0j, dtype=numpy.complex64)
         expected_error_message = "waveform_data_array must be numpy.ndarray of dtype=complex128, is complex64"
-        import pytest
         with nifake.Session('dev1') as session:
             with pytest.raises(TypeError) as exc_info:
                 session.write_waveform_numpy_complex128(invalid_waveform_data)
@@ -858,7 +874,6 @@ def test_session_write_waveform_numpy_complex128_invalid_dtype(self):
     def test_session_write_waveform_numpy_complex_interleaved_i16_invalid_dtype(self):
         invalid_waveform_data = numpy.full(10, 1.0 + 1.0j, dtype=numpy.complex64)
         expected_error_message = "waveform_data_array must be numpy.ndarray of dtype=int16, is complex64"
-        import pytest
         with nifake.Session('dev1') as session:
             with pytest.raises(TypeError) as exc_info:
                 session.write_waveform_numpy_complex_interleaved_i16(invalid_waveform_data)
diff --git a/generated/nifgen/nifgen/_visatype.py b/generated/nifgen/nifgen/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nimodinst/nimodinst/_visatype.py b/generated/nimodinst/nimodinst/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nirfsg/nirfsg/_visatype.py b/generated/nirfsg/nirfsg/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/niscope/niscope/_visatype.py b/generated/niscope/niscope/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nise/nise/_visatype.py b/generated/nise/nise/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/niswitch/niswitch/_visatype.py b/generated/niswitch/niswitch/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/generated/nitclk/nitclk/_visatype.py b/generated/nitclk/nitclk/_visatype.py
@@ -14,6 +14,7 @@
 ViInt32 = ctypes.c_int32
 ViUInt32 = ctypes.c_uint32
 ViInt64 = ctypes.c_int64
+ViUInt64 = ctypes.c_uint64
 ViString = ctypes.c_char_p
 ViReal32 = ctypes.c_float
 ViReal64 = ctypes.c_double
diff --git a/src/nifake/metadata/enums.py b/src/nifake/metadata/enums.py
@@ -239,6 +239,33 @@
             }
         ]
     },
+    'IntFlagEnum': {
+        'codegen_method': 'public',
+        'enum_class': 'IntFlag',
+        'values': [
+            {
+                'documentation': {
+                    'description': 'Flag A option.'
+                },
+                'name': 'NIFAKE_VAL_FLAG_A',
+                'value': 1
+            },
+            {
+                'documentation': {
+                    'description': 'Flag B option.'
+                },
+                'name': 'NIFAKE_VAL_FLAG_B',
+                'value': 2
+            },
+            {
+                'documentation': {
+                    'description': 'Flag C option.'
+                },
+                'name': 'NIFAKE_VAL_FLAG_C',
+                'value': 9223372036854775808
+            }
+        ]
+    },
     'MobileOSNames': {
         'values': [
             {
diff --git a/src/nifake/metadata/functions.py b/src/nifake/metadata/functions.py
@@ -348,6 +348,34 @@
         ],
         'returns': 'ViStatus'
     },
+    'FunctionWithIntflagParameter': {
+        'codegen_method': 'public',
+        'documentation': {
+            'description': "Calls a function that takes a flag parameter which can be OR'd from multiple enum values."
+        },
+        'included_in_proto': True,
+        'parameters': [
+            {
+                'direction': 'in',
+                'documentation': {
+                    'description': 'Identifies a particular instrument session.'
+                },
+                'name': 'vi',
+                'type': 'ViSession'
+            },
+            {
+                'direction': 'in',
+                'documentation': {
+                    'description': 'A flag parameter that can be a combination (bitwise OR) of IntFlagEnum values.'
+                },
+                'enum': 'IntFlagEnum',
+                'grpc_enum': None,
+                'name': 'flag',
+                'type': 'ViUInt64'
+            }
+        ],
+        'returns': 'ViStatus'
+    },
     'FunctionWithRepeatedCapabilityType': {
         'documentation': {
             'description': 'A function with a parameter that specifies repeated_capability_type.'
diff --git a/src/nifake/unit_tests/test_session.py b/src/nifake/unit_tests/test_session.py
