Resolve review comments

This commit is contained in:
Artsiom Musin
2022-11-09 15:54:16 +01:00
parent 9b047c2af6
commit c39172f516
3 changed files with 29 additions and 31 deletions

View File

@@ -213,13 +213,23 @@ class ApiV2(base.Base):
assets = (self._export(asset, post_fields) for asset in endpoint.results)
return [asset for asset in assets if asset is not None]
def _check_for_int(self, value):
return isinstance(value, int) or (isinstance(value, str) and value.isdecimal())
def _filtered_list(self, endpoint, value):
if isinstance(value, int) or value.isdecimal():
if isinstance(value, list) and len(value) == 1:
value = value[0]
if self._check_for_int(value):
return endpoint.get(id=int(value))
options = self._cache.get_options(endpoint)
identifier = next(field for field in options['search_fields'] if field in ('name', 'username', 'hostname'))
if len(value.split(',')) > 0:
identifier += '__in'
if isinstance(value, list):
if all(self._check_for_int(item) for item in value):
identifier = 'or__id'
else:
identifier = 'or__' + identifier
return endpoint.get(**{identifier: value}, all_pages=True)
def export_assets(self, **kwargs):

View File

@@ -161,7 +161,7 @@ class Export(CustomCommand):
# 1) the resource flag is not used at all, which will result in the attr being None
# 2) the resource flag is used with no argument, which will result in the attr being ''
# 3) the resource flag is used with an argument, and the attr will be that argument's value
resources.add_argument('--{}'.format(resource), nargs='?', const='')
resources.add_argument('--{}'.format(resource), nargs='*')
def handle(self, client, parser):
self.extend_parser(parser)