Add create method to BadgeCondition

Author: ErnestoLoma

pyatlan/cache/custom_metadata_cache.py

@@ -91,7 +91,7 @@ def refresh_cache(cls) -> None:
                         cls.types_by_asset[asset_type].add(attrib_type)
 
     @classmethod
-    def get_id_for_name(cls, name: str) -> Optional[str]:
+    def get_id_for_name(cls, name: str) -> str:
         """
         Translate the provided human-readable custom metadata set name to its Atlan-internal ID string.
         """
@@ -113,7 +113,7 @@ def get_id_for_name(cls, name: str) -> Optional[str]:
         )
 
     @classmethod
-    def get_name_for_id(cls, idstr: str) -> Optional[str]:
+    def get_name_for_id(cls, idstr: str) -> str:
         """
         Translate the provided Atlan-internal custom metadata ID string to the human-readable custom metadata set name.
         """
@@ -174,23 +174,32 @@ def get_all_custom_attributes(
         return m
 
     @classmethod
-    def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> Optional[str]:
+    def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> str:
         """
         Translate the provided human-readable custom metadata set and attribute names to the Atlan-internal ID string
         for the attribute.
         """
-        attr_id = None
-        if set_id := cls.get_id_for_name(set_name):
-            if sub_map := cls.map_attr_name_to_id.get(set_id):
-                attr_id = sub_map.get(attr_name)
+        set_id = cls.get_id_for_name(set_name)
+        if sub_map := cls.map_attr_name_to_id.get(set_id):
+            attr_id = sub_map.get(attr_name)
             if attr_id:
                 # If found, return straight away
                 return attr_id
-            # Otherwise, refresh the cache and look again (could be stale)
-            cls.refresh_cache()
-            if sub_map := cls.map_attr_name_to_id.get(set_id):
-                return sub_map.get(attr_name)
-        return None
+        # Otherwise, refresh the cache and look again (could be stale)
+        cls.refresh_cache()
+        if sub_map := cls.map_attr_name_to_id.get(set_id):
+            attr_id = sub_map.get(attr_name)
+            if attr_id:
+                # If found, return straight away
+                return attr_id
+            raise NotFoundError(
+                message=f"Custom metadata property with name {attr_name} does not exist in custom metadata {set_name}.",
+                code="ATLAN-PYTHON-404-009",
+            )
+        raise NotFoundError(
+            message=f"Custom metadata with ID {set_id} does not exist.",
+            code="ATLAN-PYTHON-404-009",
+        )
 
     @classmethod
     def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> Optional[str]:
@@ -199,8 +208,7 @@ def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> Optional[str]:
         for the attribute.
         """
         if sub_map := cls.map_attr_id_to_name.get(set_id):
-            attr_name = sub_map.get(attr_id)
-            if attr_name:
+            if attr_name := sub_map.get(attr_id):
                 return attr_name
             cls.refresh_cache()
             if sub_map := cls.map_attr_id_to_name.get(set_id):

pyatlan/generator/templates/entity.jinja2

@@ -11,7 +11,13 @@ from typing import Any, ClassVar, Dict, List, Optional, TypeVar
 
 from pydantic import Field, StrictStr, root_validator, validator
 
-from pyatlan.model.core import Announcement, AtlanObject, CustomMetadata, Classification, Meaning
+from pyatlan.model.core import (
+    Announcement,
+    AtlanObject,
+    CustomMetadata,
+    Classification,
+    Meaning,
+)
 from pyatlan.model.enums import (
     ADLSAccessTier,
     ADLSAccountStatus,
@@ -37,7 +43,7 @@ from pyatlan.model.enums import (
     QueryUsernameStrategy,
     SourceCostUnitType,
     GoogleDatastudioAssetType,
-    PowerbiEndorsement
+    PowerbiEndorsement,
 )
 from pyatlan.model.internal import AtlasServer, Internal
 from pyatlan.model.structs import (
@@ -55,9 +61,10 @@ from pyatlan.model.structs import (
     PopularityInsights,
     SourceTagAttribute,
 )
-from pyatlan.utils import next_id
+from pyatlan.utils import next_id, validate_required_fields
 from urllib.parse import quote, unquote
 
+
 def validate_single_required_field(field_names: list[str], values: list[Any]):
     indexes = [idx for idx, value in enumerate(values) if value is not None]
     if not indexes:
@@ -70,39 +77,8 @@ def validate_single_required_field(field_names: list[str], values: list[Any]):
             f"Only one of the following parameters are allowed: {', '.join(names)}"
         )
 
-def validate_required_fields(field_names: list[str], values: list[Any]):
-    for field_name, value in zip(field_names, values):
-        if value is None:
-            raise ValueError(f"{field_name} is required")
-        if isinstance(value, str) and not value.strip():
-            raise ValueError(f"{field_name} cannot be blank")
-        if isinstance(value, list) and len(value) == 0:
-            raise ValueError(f"{field_name} cannot be an empty list")
-{%- macro gen_properties(attribute_defs, additional_names=[]) %}
-    _convience_properties: ClassVar[list[str]] = [
-        {%- for attribute_def in attribute_defs %}
-        "{{ attribute_def.name | to_snake_case }}",
-        {%- endfor %}
-        {%- for name in additional_names  %}
-        "{{ name }}",
-         {%- endfor %}]
-
-    {%- for attribute_def in attribute_defs %}
-    {%- set type = attribute_def.typeName | get_type %}
-    {%- set property_type %}{% if attribute_def.isOptional %}Optional[{% endif %}{{type}}{% if attribute_def.isOptional %}]{% endif %}{% endset %}
-    {%- set property_name =  attribute_def.name | to_snake_case %}
-    @property
-    def {{ property_name }}(self)->{{ property_type }}:
-        return self.attributes.{{ property_name }}
-
-    @{{ property_name }}.setter
-    def {{ property_name }}(self, {{ property_name }}:{{ property_type }}):
-        if self.attributes is None:
-            self.attributes = self.Attributes()
-        self.attributes.{{ property_name }} = {{ property_name }}
 
-    {%- endfor %}
-{% endmacro %}
+{% from 'macros.jinja2' import gen_properties %}
 SelfAsset = TypeVar("SelfAsset", bound="Asset")
 {% for entity_def in entity_defs %}
 {%- set super_classes = ['AtlanObject'] if not entity_def.super_types else entity_def.super_types -%}
@@ -367,11 +343,20 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
     {%- if entity_def.name == "Badge" %}
     @classmethod
     # @validate_arguments()
-    def create(cls, *, name: StrictStr, cm_name: str, cm_attribute: str) -> Badge:
-        validate_required_fields(["name", "cm_name", "cm_attribute"], [name])
+    def create(
+        cls,
+        *,
+        name: StrictStr,
+        cm_name: str,
+        cm_attribute: str,
+        badge_conditions: list[BadgeCondition],
+    ) -> Badge:
         return cls(
             attributes=Badge.Attributes.create(
-                name=name, cm_name=cm_name, cm_attribute=cm_attribute
+                name=name,
+                cm_name=cm_name,
+                cm_attribute=cm_attribute,
+                badge_conditions=badge_conditions,
             )
         )
     {%- elif entity_def.name == "Readme" %}
@@ -540,10 +525,16 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
         @classmethod
         # @validate_arguments()
         def create(
-            cls, *, name: StrictStr, cm_name: str, cm_attribute: str
+            cls,
+            *,
+            name: StrictStr,
+            cm_name: str,
+            cm_attribute: str,
+            badge_conditions: list[BadgeCondition],
         ) -> Badge.Attributes:
             validate_required_fields(
-                ["name", "cm_name", "cm_attribute"], [name, cm_name, cm_attribute]
+                ["name", "cm_name", "cm_attribute", "badge_conditions"],
+                [name, cm_name, cm_attribute, badge_conditions],
             )
             from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 
@@ -555,6 +546,7 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 name=name,
                 qualified_name=f"badges/global/{cm_id}.{cm_attr_id}",
                 badge_metadata_attribute=f"{cm_id}.{cm_attr_id}",
+                badge_conditions=badge_conditions,
             )
         {%- elif entity_def.name == "Readme"  %}
         @classmethod

pyatlan/generator/templates/macros.jinja2

@@ -0,0 +1,25 @@
+{%- macro gen_properties(attribute_defs, additional_names=[]) %}
+    _convience_properties: ClassVar[list[str]] = [
+        {%- for attribute_def in attribute_defs %}
+        "{{ attribute_def.name | to_snake_case }}",
+        {%- endfor %}
+        {%- for name in additional_names  %}
+        "{{ name }}",
+         {%- endfor %}]
+
+    {%- for attribute_def in attribute_defs %}
+    {%- set type = attribute_def.typeName | get_type %}
+    {%- set property_type %}{% if attribute_def.isOptional %}Optional[{% endif %}{{type}}{% if attribute_def.isOptional %}]{% endif %}{% endset %}
+    {%- set property_name =  attribute_def.name | to_snake_case %}
+    @property
+    def {{ property_name }}(self)->{{ property_type }}:
+        return self.attributes.{{ property_name }}
+
+    @{{ property_name }}.setter
+    def {{ property_name }}(self, {{ property_name }}:{{ property_type }}):
+        if self.attributes is None:
+            self.attributes = self.Attributes()
+        self.attributes.{{ property_name }} = {{ property_name }}
+
+    {%- endfor %}
+{% endmacro %}

pyatlan/generator/templates/structs.jinja2

@@ -1,19 +1,48 @@
-# SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
 from datetime import datetime
-from typing import Optional
+from typing import Optional, Union
 
 from pydantic import Field
 
+from pyatlan.utils import validate_required_fields
 from pyatlan.model.core import AtlanObject
-from pyatlan.model.enums import SourceCostUnitType
+from pyatlan.model.enums import (
+    BadgeConditionColor,
+    BadgeComparisonOperator,
+    SourceCostUnitType,
+    )
 
 {% for struct in struct_defs %}
 class {{struct.name}}(AtlanObject):
     """Description"""
-    class Attributes(AtlanObject):
-        {%- for attribute_def in struct.attribute_defs %}
-        {%- set type = attribute_def.type_name | get_type %}
-        {{attribute_def.name | to_snake_case }}: {% if attribute_def.is_optional %}Optional[{% endif %}{{type}}{% if attribute_def.is_optional %}]{% endif %} = Field(None, description='' , alias='{{attribute_def.name}}')
-        {%- endfor %}
+    {%- if struct.name == "BadgeCondition" %}
+    @classmethod
+    # @validate_arguments()
+    def create(
+        cls,
+        *,
+        badge_condition_operator: BadgeComparisonOperator,
+        badge_condition_value: str,
+        badge_condition_colorhex: Union[BadgeConditionColor, str]
+    ) -> "BadgeCondition":
+        validate_required_fields(
+            [
+                "badge_condition_operator",
+                "badge_condition_value",
+                "badge_condition_colorhex",
+            ],
+            [badge_condition_operator, badge_condition_value, badge_condition_colorhex],
+        )
+        return cls(
+            badge_condition_operator=badge_condition_operator.value,
+            badge_condition_value=badge_condition_value,
+            badge_condition_colorhex=badge_condition_colorhex.value
+            if isinstance(badge_condition_colorhex, BadgeConditionColor)
+            else badge_condition_colorhex,
+        )
+    {% endif %}
+    {%- for attribute_def in struct.attribute_defs %}
+    {%- set type = attribute_def.type_name | get_type %}
+    {{attribute_def.name | to_snake_case }}: {% if attribute_def.is_optional %}Optional[{% endif %}{{type}}{% if attribute_def.is_optional %}]{% endif %} = Field(None, description='' , alias='{{attribute_def.name}}')
+    {%- endfor %}
 {% endfor %}

pyatlan/model/assets.py

@@ -62,7 +62,7 @@
     PopularityInsights,
     SourceTagAttribute,
 )
-from pyatlan.utils import next_id
+from pyatlan.utils import next_id, validate_required_fields
 
 
 def validate_single_required_field(field_names: list[str], values: list[Any]):
@@ -78,16 +78,6 @@ def validate_single_required_field(field_names: list[str], values: list[Any]):
         )
 
 
-def validate_required_fields(field_names: list[str], values: list[Any]):
-    for field_name, value in zip(field_names, values):
-        if value is None:
-            raise ValueError(f"{field_name} is required")
-        if isinstance(value, str) and not value.strip():
-            raise ValueError(f"{field_name} cannot be blank")
-        if isinstance(value, list) and len(value) == 0:
-            raise ValueError(f"{field_name} cannot be an empty list")
-
-
 SelfAsset = TypeVar("SelfAsset", bound="Asset")
 
 
@@ -3630,11 +3620,20 @@ def validate_type_name(cls, v):
 
     @classmethod
     # @validate_arguments()
-    def create(cls, *, name: StrictStr, cm_name: str, cm_attribute: str) -> Badge:
-        validate_required_fields(["name", "cm_name", "cm_attribute"], [name])
+    def create(
+        cls,
+        *,
+        name: StrictStr,
+        cm_name: str,
+        cm_attribute: str,
+        badge_conditions: list[BadgeCondition],
+    ) -> Badge:
         return cls(
             attributes=Badge.Attributes.create(
-                name=name, cm_name=cm_name, cm_attribute=cm_attribute
+                name=name,
+                cm_name=cm_name,
+                cm_attribute=cm_attribute,
+                badge_conditions=badge_conditions,
             )
         )
 
@@ -3667,10 +3666,16 @@ class Attributes(Asset.Attributes):
         @classmethod
         # @validate_arguments()
         def create(
-            cls, *, name: StrictStr, cm_name: str, cm_attribute: str
+            cls,
+            *,
+            name: StrictStr,
+            cm_name: str,
+            cm_attribute: str,
+            badge_conditions: list[BadgeCondition],
         ) -> Badge.Attributes:
             validate_required_fields(
-                ["name", "cm_name", "cm_attribute"], [name, cm_name, cm_attribute]
+                ["name", "cm_name", "cm_attribute", "badge_conditions"],
+                [name, cm_name, cm_attribute, badge_conditions],
             )
             from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 
@@ -3682,6 +3687,7 @@ def create(
                 name=name,
                 qualified_name=f"badges/global/{cm_id}.{cm_attr_id}",
                 badge_metadata_attribute=f"{cm_id}.{cm_attr_id}",
+                badge_conditions=badge_conditions,
             )
 
     attributes: "Badge.Attributes" = Field(

pyatlan/model/enums.py

@@ -262,3 +262,19 @@ class LineageDirection(str, Enum):
     UPSTREAM = "INPUT"
     DOWNSTREAM = "OUTPUT"
     BOTH = "BOTH"
+
+
+class BadgeComparisonOperator(str, Enum):
+    GT = "gt"
+    GTE = "gte"
+    LT = "lt"
+    LTE = "lte"
+    EQ = "eq"
+    NEQ = "neq"
+
+
+class BadgeConditionColor(str, Enum):
+    GREEN = "#047960"
+    YELLOW = "#F7B43D"
+    RED = "#BF1B1B"
+    GREY = "#525C73"