APP-4522

ThreatConnect-Inc · Jun 27, 2024 · 22b62df · 22b62df
1 parent 821e83d
commit 22b62df
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 34 deletions.
diff --git a/tcex/api/tc/v3/_gen/_gen_object_abc.py b/tcex/api/tc/v3/_gen/_gen_object_abc.py
@@ -471,36 +471,20 @@ def replace_artifact(self, **kwargs):
                 f'''):'''
             ),
             f'''{self.i2}"""Replace {type_.singular()} on the object."""''',
-            f'''{self.i2}if isinstance(data, ObjectABC):''',
-            f'''{self.i3}transformed_data = data''',
-            (
-                f'''{self.i2}elif isinstance(data, list) and '''
-                f'''all(isinstance(item, {model_class}) for item in data):'''
-            ),
-            f'''{self.i3}transformed_data = data''',
-            (
-                f'''{self.i2}elif isinstance(data, list) and '''
-                f'''all(isinstance(item, ObjectABC) for item in data):'''
-            ),
+            f'''{self.i2}if not isinstance(data, list):''',
+            f'''{self.i3}data = [data]''',
+            '',
+            f'''{self.i2}if all(isinstance(item, ({model_class}, ObjectABC)) for item in data):'''
             f'''{self.i3}transformed_data = data''',
-            (
-                f'''{self.i2}elif isinstance(data, list) and '''
-                f'''all(isinstance(item, dict) for item in data):'''
-            ),
+            f'''{self.i2}elif all(isinstance(item, dict) for item in data):'''
             f'''{self.i3}transformed_data = [{model_class}(**d) for d in data]''',
-            f'''{self.i2}elif isinstance(data, dict):''',
-            f'''{self.i3}transformed_data = {model_class}(**data)''',
-            f'''{self.i2}elif isinstance(data, {model_class}):''',
-            f'''{self.i3}transformed_data = data''',
             f'''{self.i2}else:'''
             f'''{self.i3}raise ValueError("Invalid data to replace_{model_type.singular()}")''',
             '',
             '',
-            f'''{self.i2}if isinstance(transformed_data, list):''',
-            f'''{self.i3}for item in transformed_data:''',
-            f'''{self.i4}item._staged = True''',
-            f'''{self.i2}elif isinstance(transformed_data, {model_class}):''',
-            f'''{self.i3}transformed_data._staged = True''',
+            f'''{self.i2}for item in transformed_data:''',
+            f'''{self.i3}item._staged = True''',
+            '',
             f'''{self.i2}self.model.{model_reference} = transformed_data  # type: ignore''',
             '',
             '',

diff --git a/tcex/api/tc/v3/intel_requirements/intel_requirement.py b/tcex/api/tc/v3/intel_requirements/intel_requirement.py
@@ -182,20 +182,16 @@ def replace_keyword_section(self, data: dict | list | ObjectABC | KeywordSection
         if not isinstance(data, list):
             data = [data]
 
-        if isinstance(data, list) and all(
-            isinstance(item, (ObjectABC, KeywordSectionModel)) for item in data
-        ):
+        if all(isinstance(item, (KeywordSectionModel, ObjectABC)) for item in data):
             transformed_data = data
-        elif isinstance(data, list) and all(isinstance(item, dict) for item in data):
+        elif all(isinstance(item, dict) for item in data):
             transformed_data = [KeywordSectionModel(**d) for d in data]
-        elif isinstance(data, dict):
-            transformed_data = KeywordSectionModel(**data)
         else:
-            raise ValueError('Invalid data to replace_keyword_section')
+            raise ValueError("Invalid data to replace_keyword_section")
+
+        for item in transformed_data:
+            item._staged = True
 
-        if isinstance(transformed_data, list):
-            for item in transformed_data:
-                item._staged = True
         self.model.keyword_sections = transformed_data  # type: ignore
 
     def stage_tag(self, data: dict | ObjectABC | TagModel):