feat(vmware_custom_attribute): Enhance custom attribute management – extended object types, refactored code

Author: sbaerlocher

plugins/modules/vmware_custom_attribute.py

@@ -1,23 +1,39 @@
 #!/usr/bin/python
 # -*- coding: utf-8 -*-
-
 # Copyright: (c) 2022, Ansible Project
 # GNU General Public License v3.0+ (see LICENSES/GPL-3.0-or-later.txt or https://www.gnu.org/licenses/gpl-3.0.txt)
 # SPDX-License-Identifier: GPL-3.0-or-later
+"""
+Module for managing VMware custom attribute definitions.
+"""
+
+from types import ModuleType
+from typing import Any, Dict, Optional
 
-from __future__ import absolute_import, division, print_function
-__metaclass__ = type
+from ansible.module_utils.basic import AnsibleModule
+from ansible_collections.community.vmware.plugins.module_utils.vmware import (
+    PyVmomi,
+    vmware_argument_spec,
+)
 
+try:
+    from pyVmomi import vim as _vim  # type: ignore
+except ImportError:
+    _vim = None  # type: ignore
 
-DOCUMENTATION = r'''
+vim: Optional[ModuleType] = _vim
+
+DOCUMENTATION = r"""
 ---
 module: vmware_custom_attribute
-version_added: '3.2.0'
-short_description: Manage custom attributes definitions
+version_added: "3.3.0"
+short_description: Manage custom attribute definitions for vSphere objects.
 description:
-  - This module can be used to add and remove custom attributes definitions for various vSphere objects.
+  - This module adds or removes custom attribute definitions for various vSphere objects.
+  - It supports all object types provided by VMware (e.g. Cluster, Datacenter, VirtualMachine, etc.).
 author:
   - Mario Lenz (@mariolenz)
+  - Simon Baerlocher (@sbaerlocher)
 options:
   custom_attribute:
     description:
@@ -27,6 +43,8 @@
   object_type:
     description:
       - Type of the object the custom attribute is associated with.
+      - All supported types are listed here.
+    required: true
     type: str
     choices:
       - Cluster
@@ -39,7 +57,8 @@
       - HostSystem
       - ResourcePool
       - VirtualMachine
-    required: true
+      - Network
+      - VirtualApp
   state:
     description:
       - Manage definition of custom attributes.
@@ -51,11 +70,10 @@
     choices: ['present', 'absent']
     type: str
 extends_documentation_fragment:
-- community.vmware.vmware.documentation
-
-'''
+  - community.vmware.vmware.documentation
+"""
 
-EXAMPLES = r'''
+EXAMPLES = r"""
 - name: Add VM Custom Attribute Definition
   community.vmware.vmware_custom_attribute:
     hostname: "{{ vcenter_hostname }}"
@@ -65,7 +83,7 @@
     object_type: VirtualMachine
     custom_attribute: custom_attr_def_1
   delegate_to: localhost
-  register: defs
+  register: definitions
 
 - name: Remove VM Custom Attribute Definition
   community.vmware.vmware_custom_attribute:
@@ -76,100 +94,182 @@
     object_type: VirtualMachine
     custom_attribute: custom_attr_def_1
   delegate_to: localhost
-  register: defs
-'''
+  register: definitions
 
-RETURN = r'''
-'''
+- name: Add Network Custom Attribute Definition
+  community.vmware.vmware_custom_attribute:
+    hostname: "{{ vcenter_hostname }}"
+    username: "{{ vcenter_username }}"
+    password: "{{ vcenter_password }}"
+    state: present
+    object_type: Network
+    custom_attribute: custom_attr_network
+  delegate_to: localhost
+  register: definitions
+"""
 
-from ansible.module_utils.basic import AnsibleModule
-from ansible_collections.community.vmware.plugins.module_utils.vmware import PyVmomi, vmware_argument_spec
+RETURN = r"""
+changed:
+    description: Indicates if any change was made.
+    type: bool
+failed:
+    description: Indicates if the operation failed.
+    type: bool
+"""
 
-try:
-    from pyVmomi import vim
-except ImportError:
-    pass
 
+def get_object_type_mapping() -> Dict[str, Any]:
+    """Returns a mapping from object type names to the corresponding pyVmomi classes."""
+    return {
+        "Cluster": vim.ClusterComputeResource if vim else None,
+        "Datacenter": vim.Datacenter if vim else None,
+        "Datastore": vim.Datastore if vim else None,
+        "DistributedVirtualPortgroup": (
+            vim.dvs.DistributedVirtualPortgroup if vim else None
+        ),
+        "DistributedVirtualSwitch": vim.DistributedVirtualSwitch if vim else None,
+        "Folder": vim.Folder if vim else None,
+        "Global": None,
+        "HostSystem": vim.HostSystem if vim else None,
+        "ResourcePool": vim.ResourcePool if vim else None,
+        "VirtualMachine": vim.VirtualMachine if vim else None,
+        "Network": vim.Network if vim else None,
+        "VirtualApp": getattr(vim, "VirtualApp", None) if vim else None,
+    }
 
-class CustomAttribute(PyVmomi):
-    def __init__(self, module):
-        super(CustomAttribute, self).__init__(module)
+
+class CustomAttributeManager(PyVmomi):
+    """Class responsible for managing custom attribute definitions."""
+
+    def __init__(self, module: AnsibleModule) -> None:
+        super().__init__(module)
+        self.module = module
+
+        if not isinstance(module.params, dict):
+            self.module.fail_json(msg="module.params is not a dict")
+        self.parameters: Dict[str, Any] = module.params
+
+        custom_attribute_value = self.parameters.get("custom_attribute", "")
+        if (
+            not isinstance(custom_attribute_value, str)
+            or not custom_attribute_value.strip()
+        ):
+            self.module.fail_json(msg="'custom_attribute' must be a non-empty string")
+
+        if vim is None:
+            self.module.fail_json(msg="pyVmomi is required for this module")
 
         if not self.is_vcenter():
-            self.module.fail_json(msg="You have to connect to a vCenter server!")
-
-        object_types_map = {
-            'Cluster': vim.ClusterComputeResource,
-            'Datacenter': vim.Datacenter,
-            'Datastore': vim.Datastore,
-            'DistributedVirtualPortgroup': vim.DistributedVirtualPortgroup,
-            'DistributedVirtualSwitch': vim.DistributedVirtualSwitch,
-            'Folder': vim.Folder,
-            'Global': None,
-            'HostSystem': vim.HostSystem,
-            'ResourcePool': vim.ResourcePool,
-            'VirtualMachine': vim.VirtualMachine
-        }
-
-        self.object_type = object_types_map[self.params['object_type']]
-
-    def remove_custom_def(self, field):
-        changed = False
-        for x in self.custom_field_mgr:
-            if x.name == field and x.managedObjectType == self.object_type:
-                changed = True
-                if not self.module.check_mode:
-                    self.content.customFieldsManager.RemoveCustomFieldDef(key=x.key)
-                break
-        return {'changed': changed, 'failed': False}
-
-    def add_custom_def(self, field):
-        changed = False
-        found = False
-        for x in self.custom_field_mgr:
-            if x.name == field and x.managedObjectType == self.object_type:
-                found = True
-                break
-
-        if not found:
-            changed = True
+            self.module.fail_json(msg="A connection to a vCenter server is required!")
+
+        object_type_value = self.parameters.get("object_type", "")
+        if not isinstance(object_type_value, str) or not object_type_value.strip():
+            self.module.fail_json(msg="'object_type' must be a non-empty string")
+
+        object_type_mapping = get_object_type_mapping()
+        self.object_type = object_type_mapping.get(object_type_value)
+        if self.object_type is None and object_type_value != "Global":
+            self.module.fail_json(msg=f"Unsupported object type: {object_type_value}")
+
+        try:
+            self.custom_field_definitions = self.content.customFieldsManager.field
+        except AttributeError:
+            self.module.fail_json(
+                msg="Failed to access customFieldsManager in vCenter content"
+            )
+
+    def find_custom_attribute_definition(
+        self, custom_attribute_name: str
+    ) -> Optional[Any]:
+        """Searches for a custom attribute definition and returns it if found."""
+        for custom_field_definition in self.custom_field_definitions:
+            if (
+                custom_field_definition.name == custom_attribute_name
+                and custom_field_definition.managedObjectType == self.object_type
+            ):
+                return custom_field_definition
+        return None
+
+    def remove_custom_definition(self, custom_attribute_name: str) -> Dict[str, Any]:
+        """Removes the custom attribute definition if it exists."""
+        state_changed = False
+        custom_field_definition = self.find_custom_attribute_definition(
+            custom_attribute_name
+        )
+        if custom_field_definition:
+            state_changed = True
+            if not self.module.check_mode:
+                self.content.customFieldsManager.RemoveCustomFieldDef(
+                    key=custom_field_definition.key
+                )
+        return {"changed": state_changed, "failed": False}
+
+    def add_custom_definition(self, custom_attribute_name: str) -> Dict[str, Any]:
+        """Adds the custom attribute definition if it does not exist."""
+        state_changed = False
+        if not self.find_custom_attribute_definition(custom_attribute_name):
+            state_changed = True
             if not self.module.check_mode:
-                self.content.customFieldsManager.AddFieldDefinition(name=field, moType=self.object_type)
-        return {'changed': changed, 'failed': False}
-
-
-def main():
-    argument_spec = vmware_argument_spec()
-    argument_spec.update(
-        custom_attribute=dict(type='str', no_log=False, required=True),
-        object_type=dict(type='str', required=True, choices=[
-            'Cluster',
-            'Datacenter',
-            'Datastore',
-            'DistributedVirtualPortgroup',
-            'DistributedVirtualSwitch',
-            'Folder',
-            'Global',
-            'HostSystem',
-            'ResourcePool',
-            'VirtualMachine'
-        ]),
-        state=dict(type='str', default='present', choices=['absent', 'present']),
+                self.content.customFieldsManager.AddFieldDefinition(
+                    name=custom_attribute_name, moType=self.object_type
+                )
+        return {"changed": state_changed, "failed": False}
+
+
+def manage_custom_attribute_definition(module: AnsibleModule) -> None:
+    """Determines whether to add or remove the custom attribute definition based on the 'state' parameter."""
+    if not isinstance(module.params, dict):
+        module.fail_json(msg="module.params is not a dict")
+    parameters: Dict[str, Any] = module.params
+    custom_attribute_name = parameters["custom_attribute"]
+    desired_state = parameters["state"]
+    custom_attribute_manager = CustomAttributeManager(module)
+    if desired_state == "present":
+        result = custom_attribute_manager.add_custom_definition(custom_attribute_name)
+    else:
+        result = custom_attribute_manager.remove_custom_definition(
+            custom_attribute_name
+        )
+    module.exit_json(**result)
+
+
+def main() -> None:
+    """Main entry point for the module."""
+    argument_specification = vmware_argument_spec()
+    argument_specification.update(
+        custom_attribute={"type": "str", "no_log": False, "required": True},
+        object_type={
+            "type": "str",
+            "required": True,
+            "choices": [
+                "Cluster",
+                "Datacenter",
+                "Datastore",
+                "DistributedVirtualPortgroup",
+                "DistributedVirtualSwitch",
+                "Folder",
+                "Global",
+                "HostSystem",
+                "ResourcePool",
+                "VirtualMachine",
+                "Network",
+                "VirtualApp",
+            ],
+        },
+        state={"type": "str", "default": "present", "choices": ["absent", "present"]},
     )
     module = AnsibleModule(
-        argument_spec=argument_spec,
+        argument_spec=argument_specification,
         supports_check_mode=True,
     )
 
-    pyv = CustomAttribute(module)
-    results = dict(changed=False, custom_attribute_defs=list())
-    if module.params['state'] == "present":
-        results = pyv.add_custom_def(module.params['custom_attribute'])
-    elif module.params['state'] == "absent":
-        results = pyv.remove_custom_def(module.params['custom_attribute'])
-
-    module.exit_json(**results)
+    try:
+        manage_custom_attribute_definition(module)
+    except ValueError as error:
+        module.fail_json(msg=f"ValueError: {error}")
+    except KeyError as error:
+        module.fail_json(msg=f"KeyError: {error}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()