perf: del pandas

pull/8522/head
feng626 2022-06-30 17:06:50 +08:00 committed by Jiangjie.Bai
parent 011e9ffec4
commit af2d927c1f
2 changed files with 25 additions and 24 deletions

View File

@ -1,6 +1,6 @@
import os import os
import time import time
import pandas as pd from openpyxl import Workbook
from collections import defaultdict, OrderedDict from collections import defaultdict, OrderedDict
from django.conf import settings from django.conf import settings
@ -48,7 +48,7 @@ class BaseAccountHandler:
_fields = cls.get_header_fields(v) _fields = cls.get_header_fields(v)
header_fields.update(_fields) header_fields.update(_fields)
else: else:
header_fields[field] = v.label header_fields[field] = str(v.label)
return header_fields return header_fields
@classmethod @classmethod
@ -59,7 +59,7 @@ class BaseAccountHandler:
data = cls.unpack_data(serializer.data) data = cls.unpack_data(serializer.data)
row_dict = {} row_dict = {}
for field, header_name in header_fields.items(): for field, header_name in header_fields.items():
row_dict[header_name] = data[field] row_dict[header_name] = str(data[field])
return row_dict return row_dict
@ -72,24 +72,24 @@ class AssetAccountHandler(BaseAccountHandler):
return filename return filename
@classmethod @classmethod
def create_df(cls): def create_data_map(cls):
df_dict = defaultdict(list) data_map = defaultdict(list)
sheet_name = AuthBook._meta.verbose_name sheet_name = AuthBook._meta.verbose_name
accounts = AuthBook.get_queryset().select_related('systemuser') accounts = AuthBook.get_queryset().select_related('systemuser')
if not accounts.first(): if not accounts.first():
return df_dict return data_map
header_fields = cls.get_header_fields(AccountSecretSerializer(accounts.first())) header_fields = cls.get_header_fields(AccountSecretSerializer(accounts.first()))
for account in accounts: for account in accounts:
account.load_auth() account.load_auth()
row = cls.create_row(account, AccountSecretSerializer, header_fields) row = cls.create_row(account, AccountSecretSerializer, header_fields)
df_dict[sheet_name].append(row) if sheet_name not in data_map:
for k, v in df_dict.items(): data_map[sheet_name].append(list(row.keys()))
df_dict[k] = pd.DataFrame(v) data_map[sheet_name].append(list(row.values()))
logger.info('\n\033[33m- 共收集 {} 条资产账号\033[0m'.format(accounts.count())) logger.info('\n\033[33m- 共收集 {} 条资产账号\033[0m'.format(accounts.count()))
return df_dict return data_map
class AppAccountHandler(BaseAccountHandler): class AppAccountHandler(BaseAccountHandler):
@ -101,19 +101,19 @@ class AppAccountHandler(BaseAccountHandler):
return filename return filename
@classmethod @classmethod
def create_df(cls): def create_data_map(cls):
df_dict = defaultdict(list) data_map = defaultdict(list)
accounts = Account.get_queryset().select_related('systemuser') accounts = Account.get_queryset().select_related('systemuser')
for account in accounts: for account in accounts:
account.load_auth() account.load_auth()
app_type = account.type app_type = account.type
sheet_name = AppType.get_label(app_type) sheet_name = AppType.get_label(app_type)
row = cls.create_row(account, AppAccountSecretSerializer) row = cls.create_row(account, AppAccountSecretSerializer)
df_dict[sheet_name].append(row) if sheet_name not in data_map:
for k, v in df_dict.items(): data_map[sheet_name].append(list(row.keys()))
df_dict[k] = pd.DataFrame(v) data_map[sheet_name].append(list(row.values()))
logger.info('\n\033[33m- 共收集{}条应用账号\033[0m'.format(accounts.count())) logger.info('\n\033[33m- 共收集{}条应用账号\033[0m'.format(accounts.count()))
return df_dict return data_map
handler_map = { handler_map = {
@ -142,15 +142,18 @@ class AccountBackupHandler:
if not handler: if not handler:
continue continue
df_dict = handler.create_df() data_map = handler.create_data_map()
if not df_dict: if not data_map:
continue continue
filename = handler.get_filename(self.plan_name) filename = handler.get_filename(self.plan_name)
with pd.ExcelWriter(filename) as w:
for sheet, df in df_dict.items(): wb = Workbook(filename)
sheet = sheet.replace(' ', '-') for sheet, data in data_map.items():
getattr(df, 'to_excel')(w, sheet_name=sheet, index=False) ws = wb.create_sheet(str(sheet))
for row in data:
ws.append(row)
wb.save(filename)
files.append(filename) files.append(filename)
timedelta = round((time.time() - time_start), 2) timedelta = round((time.time() - time_start), 2)
logger.info('步骤完成: 用时 {}s'.format(timedelta)) logger.info('步骤完成: 用时 {}s'.format(timedelta))

View File

@ -57,8 +57,6 @@ html2text==2020.1.16
pyzipper==0.3.5 pyzipper==0.3.5
python3-saml==1.12.0 python3-saml==1.12.0
websocket-client==1.2.3 websocket-client==1.2.3
numpy==1.22.0
pandas==1.3.5
pyjwkest==1.4.2 pyjwkest==1.4.2
jsonfield2==4.0.0.post0 jsonfield2==4.0.0.post0
geoip2==4.5.0 geoip2==4.5.0