pyxllib 0.3.96__py3-none-any.whl → 0.3.197__py3-none-any.whl

pyxllib/file/xlsxlib.py

@@ -7,12 +7,14 @@
 """
 扩展了些自己的openpyxl工具
 """
+import copy
+
 import time
 
 from pyxllib.prog.pupil import check_install_package, run_once
 
 check_install_package('openpyxl')
-check_install_package('premailer')
+# check_install_package('premailer')
 # check_install_package('xlrd2')
 check_install_package('yattag')
 check_install_package('jsonpickle')
@@ -30,11 +32,13 @@ import io
 
 import xlrd
 
+import filetype
 import openpyxl
 from openpyxl import Workbook
 from openpyxl.cell.cell import MergedCell
 from openpyxl.styles import Font, Alignment
 from openpyxl.utils.cell import get_column_letter, column_index_from_string
+import openpyxl.worksheet.formula
 import pandas as pd
 
 try:
@@ -117,7 +121,7 @@ def is_valid_excel_address(address):
         return is_valid_excel_cell(address)
 
 
-@run_once('str', debug=True)
+@run_once('str')
 def xlfmt2pyfmt_date(xl_fmt):
     """ 日期的渲染操作
 
@@ -233,6 +237,8 @@ def xl_render_value(x, xl_fmt):
     注意，遇到公式是很难计算处理的，大概率只能保持原公式显示
     因为日期用的比较多，需要时常获得真实的渲染效果，所以这里封装一个接口
 
+    对于JSA等场景，直接使用Cell.Text获取渲染值就行，不需要这里这么复杂的实现
+
     >>> xl_render_value(datetime.datetime(2020, 1, 1), 'yyyy-mm-dd')
     '2020-01-01'
     """
@@ -351,19 +357,73 @@ def convert_xls_to_xlsx(xls_file):
 
 
 def load_as_xlsx_file(file_path, keep_links=False, keep_vba=False):
+    """ 这个不能全信文件给的扩展名，需要智能判断 """
+
+    # 0 工具函数
+    @run_once()
+    def read_xlsx():
+        file = file_path
+        # 如果文件原本的后缀不是xlsx，openpyxl是读不了的，要绕个弯
+        if file.suffix[1:] not in ('xlsx', 'xlsm'):
+            with open(file_path, 'rb') as f2:
+                data = f2.read()
+            file = io.BytesIO(data)
+        try:
+            return openpyxl.load_workbook(file,
+                                          keep_links=keep_links,
+                                          keep_vba=keep_vba), ''
+        except Exception as e:
+            if isinstance(e, TimeoutError):  # 这里触发的是总的超时设定
+                raise e
+            return None, format_exception(e, 2)
+
+    @run_once()
+    def read_xls():
+        try:
+            return convert_xls_to_xlsx(file_path), ''
+        except Exception as e:
+            return None, format_exception(e, 2)
+
+    @run_once()
+    def read_csv():
+        try:
+            return convert_csv_to_xlsx(file_path), ''
+        except Exception as e:
+            return None, format_exception(e, 2)
+
+    def read_test(suffix):
+        if suffix in ('xlsx', 'xlsm', 'zip'):
+            wb, error = read_xlsx()
+        elif suffix == 'xls':
+            wb, error = read_xls()
+        elif suffix == 'csv':
+            wb, error = read_csv()
+        else:
+            wb, error = None, f'不支持的文件类型：{suffix}'
+        return wb, error
+
+    # 1 优先相信用户输入的文件名类型
     file_path = Path(file_path)
-    suffix = file_path.suffix.lower()
-    if suffix in ('.xlsx', '.xlsm'):
-        wb = openpyxl.load_workbook(file_path,
-                                    keep_links=keep_links,
-                                    keep_vba=keep_vba)
-    elif suffix == '.xls':
-        wb = convert_xls_to_xlsx(file_path)
-    elif suffix == '.csv':
-        wb = convert_csv_to_xlsx(file_path)
-    else:
-        return None
-    return wb
+    suffix = file_path.suffix.lower()[1:]
+    wb, error = read_test(suffix)
+    if wb is not None:
+        return wb, suffix
+
+    # 2 如果处理不了，则尝试用filetype判断的类型
+    suffix2 = filetype.guess(file_path)
+    suffix2 = suffix2.extension if suffix2 else ''
+    wb, _ = read_test(suffix2)
+    if wb is not None:
+        return wb, suffix2
+
+    # 3 如果还处理不了，再把其他可能的情况试一遍
+    for suffix in ('xlsx', 'xls', 'csv'):
+        wb, _ = read_test(suffix)
+        if wb is not None:
+            return wb, suffix
+
+    # 4 确实是处理不了的类型，返回报错信息
+    return None, error
 
 
 def parse_range_address(address):
@@ -466,7 +526,7 @@ def is_string_type(value):
     try:
         pd.to_datetime(value, errors='raise')
         return False
-    except (ValueError, TypeError, OverflowError):
+    except (ValueError, TypeError, OverflowError, AttributeError):
         pass
 
     # 检查是否为浮点数类型
@@ -521,10 +581,18 @@ class XlCell(openpyxl.cell.cell.Cell):  # 适用于 openpyxl.cell.cell.MergedCel
 
         TODO 这个函数还是可以看看能不能有更好的实现、提速
         """
+
+        def try_offset(x, y):
+            try:
+                return isinstance(self.offset(x, y), MergedCell)
+            except ValueError:
+                # 有可能会越界：ValueError: Row numbers must be between 1 and 1048576
+                return False
+
         _type, status = 0, {}
         if isinstance(self, MergedCell):
             _type = 1
-        elif isinstance(self.offset(1, 0), MergedCell) or isinstance(self.offset(0, 1), MergedCell):
+        elif try_offset(1, 0) or try_offset(0, 1):
             # 这里只能判断可能是合并单元格，具体是不是合并单元格，还要
             rng = self.in_range()
             status['rng'] = rng
@@ -698,7 +766,7 @@ class XlCell(openpyxl.cell.cell.Cell):  # 适用于 openpyxl.cell.cell.MergedCel
         # openpyxl的机制，如果没有配置日期格式，读取到的是默认的'mm-dd-yy'，其实在中文场景，默认格式应该是后者
         if fmt == 'mm-dd-yy':
             return 'yyyy/m/d'  # 中文的默认日期格式
-        elif fmt == 'yyyy\-mm\-dd':  # 不知道为什么会有提取到这种\的情况，先暴力替换了
+        elif fmt == r'yyyy\-mm\-dd':  # 不知道为什么会有提取到这种\的情况，先暴力替换了
             fmt = 'yyyy-mm-dd'
         return fmt
 
@@ -709,7 +777,10 @@ class XlCell(openpyxl.cell.cell.Cell):  # 适用于 openpyxl.cell.cell.MergedCel
         注意，遇到公式是很难计算处理的，大概率只能保持原公式显示
         因为日期用的比较多，需要时常获得真实的渲染效果，所以这里封装一个接口
         """
+
         x = self.value
+        if isinstance(x, openpyxl.worksheet.formula.ArrayFormula):  # 数组公式要特别渲染
+            return x.text
         xl_fmt = self.get_number_format()
         return xl_render_value(x, xl_fmt)
 
@@ -1136,7 +1207,7 @@ class XlWorksheet(openpyxl.worksheet.worksheet.Worksheet):
 
         return df
 
-    def copy_range(self, src_addr, dst_cell, *, temp_sheet=False, return_mid_result=False):
+    def copy_range(self, src_addr, dst_cell, *, temp_sheet=False, return_mode=False):
         """ 将自身cell_range区间的内容、格式，拷贝到目标dst_cell里
 
         :param str src_addr: 自身的一片单元格范围
@@ -1158,7 +1229,7 @@ class XlWorksheet(openpyxl.worksheet.worksheet.Worksheet):
         mid_result = {}
         if temp_sheet:
             ws3 = self.parent.create_sheet('__copy_range')
-            mid_result = self.copy_range(src_addr, ws3['A1'], return_mid_result=True)
+            mid_result = self.copy_range(src_addr, ws3['A1'], return_mode=True)
             ws1 = ws3
             src_addr = f'A1:{excel_addr(mid_result["n"], mid_result["m"])}'
         else:
@@ -1193,7 +1264,7 @@ class XlWorksheet(openpyxl.worksheet.worksheet.Worksheet):
         if temp_sheet:
             self.parent.remove(ws1)
 
-        if return_mid_result:
+        if return_mode:
             return mid_result
 
     def reindex_columns(self, orders):
@@ -1964,7 +2035,7 @@ class XlWorkbook(openpyxl.Workbook):
                         'cells': extract_cells_content(ws)
                     })
 
-                    if not summary['cells']:  # 如果没有数据，则大概率是数据透视表，是计算出来的，读取不到~
+                    if not summary['cells']:  # 如果没有数据，则大概率是数据透视表，是计算出来的，读取不到~ 但是JSA等场景应该有办法获得
                         summary['sheetType'] = 'PivotTable'
                         del summary['cells']
                 else:
@@ -2399,7 +2470,7 @@ def extract_workbook_summary2(file_path, *,
     """
     :param keep_links: 是否保留外部表格链接数据。如果保留，打开好像会有点问题。
     :param mode:
-        0，最原始的summary3摘要
+        0，最原始的summary2摘要
         1，添加当前工作表、单元格位置的信息
     :param kwargs: 捕捉其他参数，主要是向下兼容，其实现在并没有用
 
@@ -2411,7 +2482,12 @@ def extract_workbook_summary2(file_path, *,
     res = {}
     res['fileName'] = file_path.name
     start_time = time.time()
-    wb = load_as_xlsx_file(file_path, keep_links=keep_links, keep_vba=keep_vba)
+    wb, suffix = load_as_xlsx_file(file_path, keep_links=keep_links, keep_vba=keep_vba)
+    if wb is None:
+        res['error'] = f'Load file error。{suffix}'
+    else:
+        res['fileType'] = suffix
+
     load_time = time.time() - start_time
     if wb is None:  # 不支持的文件类型，不报错，只是返回最基本的文件名信息
         if return_mode == 1:
@@ -2459,7 +2535,7 @@ def update_raw_summary2(data):
 
     # 3 判断键值顺序
     keys = list(data.keys())
-    ref_keys = ['fileName', 'chineseContentRatio', 'nonEmptyCellRatio', 'sheetNames', 'sheets']
+    ref_keys = ['fileName', 'fileType', 'chineseContentRatio', 'nonEmptyCellRatio', 'sheetNames', 'sheets']
     if keys != ref_keys:
         data = {k: data[k] for k in ref_keys if k in data}
 
@@ -2597,30 +2673,39 @@ class WorkbookSummary3:
                     for addr, _ in row:
                         new_cells[addr] = cells[addr]
 
-            new_cells = {}
-            for rows in rows_groups:
+            total_new_cells = []
+            for rows in reversed(rows_groups):
+                new_cells = {}
                 if len(rows) < 10:
                     extract_cells_from_rows(rows)
                 else:  # 压缩中间的数据
                     # 如果评估到最终摘要可能太小，要收敛下删除的范围
                     n, m = len(rows), len(rows[0])
                     target_n = int(target_reduce_cells_num / m + 0.5)  # 本来应该删除多少行才行
-                    cur_n = n - 4 if target_n > n - 4 else target_n  # 实际删除多少行
-                    left_n = n - cur_n  # 剩余多少行
-                    b = left_n // 2
-                    a = left_n - b
-
-                    extract_cells_from_rows(rows[:a])
-                    addr = combine_addresses(rows[a][0][0], rows[-b - 1][-1][0])
-                    # new_cells[addr] = '这块区域的内容跟前面几行、后面几行的内容结构是一致的，省略显示'
-                    new_cells[addr] = '...'
-                    extract_cells_from_rows(rows[-b:])
-
-                    target_reduce_cells_num -= cur_n * m
-                    if target_reduce_cells_num <= 0:
-                        break
-
-            sheet['cells'] = new_cells
+                    if target_n <= 0:  # 如果删除的行数太少，那么就不压缩了
+                        extract_cells_from_rows(rows)
+                    else:
+                        cur_n = n - 4 if target_n > n - 4 else target_n  # 实际删除多少行
+                        left_n = n - cur_n  # 剩余多少行
+                        b = left_n // 2
+                        a = left_n - b
+
+                        extract_cells_from_rows(rows[:a])
+                        addr = combine_addresses(rows[a][0][0], rows[-b - 1][-1][0])
+                        # new_cells[addr] = '这块区域的内容跟前面几行、后面几行的内容结构是一致的，省略显示'
+                        new_cells[addr] = '...'
+                        extract_cells_from_rows(rows[-b:])
+
+                        target_reduce_cells_num -= cur_n * m
+                        # 240429周一21:57，这两行不能开，否则会过渡精简。如果压缩够了，那么后面的单元格需要全量补上。
+                        # if target_reduce_cells_num <= 0:  # 满足以后不是直接break，而是要把后续的内容都保留
+                        #     break
+                total_new_cells.append(new_cells)
+
+            new_cells2 = {}
+            for rows in reversed(total_new_cells):
+                new_cells2.update(rows)
+            sheet['cells'] = new_cells2
 
     @classmethod
     def reduce4_truncate_cells(cls, y, summary_limit_len, *, cur_summary_len=None):
@@ -2752,6 +2837,7 @@ class WorkbookSummary3:
         if cur_summary_len is None:
             cur_summary_len = cls.count_length(y)
 
+        cur_summary_len0 = cur_summary_len
         active_sheet = y['ActiveSheet']
 
         # 1 预计要删除单元格数
@@ -2783,7 +2869,8 @@ class WorkbookSummary3:
             return cls.count_length(y)
 
         # 4 否则每张表按照比例删单元格，只保留前面部分的单元格
-        left_rate = 1 - r  # 原始保留比例
+        # todo 这里应该有更好的筛选机制，后续可以思考思考
+        left_rate = min((summary_limit_len + cur_summary_len) / (2 * cur_summary_len), 0.9)  # 首轮减小一点调整幅度
         while True:
             for i, st in enumerate(y['sheets']):
                 if i == active_sheet_index:
@@ -2795,10 +2882,10 @@ class WorkbookSummary3:
             cur_summary_len = cls.count_length(y)
             if cur_summary_len <= summary_limit_len:
                 return cur_summary_len
-            if left_rate * total_cells_num < 1:
+            if left_rate * total_cells_num < 1:  # 都没有单元格，别删了
                 break
-            else:
-                left_rate *= 0.8  # 缩小保留比例，再试
+            else:  # 更新保留比率，再试
+                left_rate *= min(summary_limit_len / cur_summary_len, 0.9)
 
         return cur_summary_len
 
@@ -2881,8 +2968,12 @@ class WorkbookSummary3:
             return y
 
         x = summary2
+        if 'error' in x:
+            return x
+
         y = {
             'fileName': x['fileName'],
+            'fileType': x['fileType'],
             'sheetNames': x['sheetNames'],
             'sheets': x['sheets'],
             'mode': 'Complete information',
@@ -2927,33 +3018,90 @@ def extract_workbook_summary3(file_path, summary_limit_len=4000, **kwargs):
     return data
 
 
+def summary2_add_enums(summary2, enum_values):
+    # 1 预备
+    if enum_values is True:
+        enum_values = (20, 10)
+    max_len, max_num = enum_values
+
+    # 2 枚举值
+    for sheet in summary2['sheets']:
+        # 2.1 遍历计数
+        cols = defaultdict(Counter)
+        for addr, val in sheet['cells'].items():
+            n = len(str(val))
+            if not n or n > max_len:
+                continue
+            col = re.match(r'[A-Z]+', addr).group()
+            cols[col][val] += 1
+
+        # 2.2 添加枚举值列
+        enums = {}
+        keys = sorted(cols.keys(), key=column_index_from_string)
+        for k in keys:
+            ct = cols[k]
+            if len(ct) > max_num:
+                continue
+            vals = ct.most_common()
+            if vals[0][1] == 1:  # 都只出现了一次，也不认为是枚举值，跳过。或者是小数据表，一般也能全量展示。
+                continue
+            enums[k] = [v for v, _ in vals]
+
+        # 2.3 保存
+        if enums:
+            sheet['enums'] = enums
+            # enums2 = json.dumps(enums, ensure_ascii=False, default=str)
+            # sheet['enums'] = json.loads(enums2)
+
+    return summary2
+
+
 def extract_workbook_summary3b(file_path,
                                summary_limit_len=4000,
                                timeout_seconds=10,
                                return_mode=0,
                                debug=False,
                                len_mode=0,
+                               enum_values=False,
                                **kwargs):
     """
 
     :param summary_limit_len: 摘要长度限制
     :param timeout_seconds: 超时限制
-    :param return_mode: 返回模式，0表示只返回摘要，1表示返回摘要和耗时
+    :param return_mode: 返回模式
+        0，表示只返回摘要
+        1，表示返回摘要和耗时
+        2, 再增加返回summary2
     :param len_mode:
         0, 使用len作为token长度评估
         1, 使用模型评估实际token长度
+    :param enum_values: 是否展示每列枚举值
+        False, 默认不展示
+        True, 展示，并且默认参数 (20, 10) 表示长度超过20的丢弃，只保留枚举类型不超过10种值的列
     :param kwargs: 其他是summary2读取文件的时候的参数，其实都不太关键，一般不用特地设置
     """
     res = {}
     res['fileName'] = Path(file_path).name
     load_time = summary2_time = summary3_time = -1
+    summary2_res = {}
+
+    def reduce_summary(summary):
+        """ 如果转json后的summary超过4K，去掉可能的sheets字段 """
+        s = json.dumps(summary, ensure_ascii=False)
+        if len(s) < 4000:
+            if 'sheets' in summary:
+                del summary['sheets']
 
     try:
         with Timeout(timeout_seconds):
             start_time = time.time()
             res, load_time = extract_workbook_summary2(file_path, mode=1, return_mode=1, **kwargs)
             # res = convert_to_json_compatible(res)
+            summary2_res = copy.deepcopy(res)
+            if enum_values:
+                res = summary2_add_enums(res, enum_values)
             summary2_time = time.time() - start_time - load_time
+
             start_time = time.time()
             if len_mode == 1:
                 res = WorkbookSummary3plus.summary2_to_summary3b(res, summary_limit_len)
@@ -2964,13 +3112,20 @@ def extract_workbook_summary3b(file_path,
         if debug:
             raise e
         res['error'] = f'超时，未完成摘要提取：{timeout_seconds}秒'
+        reduce_summary(res)
     except Exception as e:
         if debug:
             raise e
         res['error'] = f'提取摘要时发生错误：{format_exception(e, 2)}'
+        reduce_summary(res)
+
+    time_dict = {'load_time': human_readable_number(load_time),
+                 'summary2_time': human_readable_number(summary2_time),
+                 'summary3_time': human_readable_number(summary3_time)}
 
     if return_mode == 1:
-        return res, {'load_time': human_readable_number(load_time),
-                     'summary2_time': human_readable_number(summary2_time),
-                     'summary3_time': human_readable_number(summary3_time)}
+        return res, time_dict
+    elif return_mode == 2:
+        return res, time_dict, summary2_res
+
     return res