Finds sheet by name, return compressed path

1 files changed, 13 insertions, 3 deletions

diff --git a/compress.py b/compress.py
index 9e92790..816a872 100755
--- a/compress.py
+++ b/compress.py
@@ -20,11 +20,19 @@ def wrap_districts(districts, length=32):
     return lines
 
 
-def compress(filepath):
+def compress(filepath: str) -> str:
     regions = {}  # all high/middle risk regions go here
     # open xlsx file
     xlsx = load_workbook(filepath)
-    lst = xlsx.worksheets[0]  # "列表"
+    lst = None
+    for sheet in xlsx.worksheets:
+        if sheet.title == "列表":
+            lst = sheet
+            break
+    else:
+        # sheet not found
+        raise Exception("未找到名为“列表”的表格")
+
     # B3:Dxxxx
     for province, city, district in lst.iter_rows(
         min_row=3, max_row=None, min_col=2, max_col=4, values_only=True
@@ -81,7 +89,9 @@ def compress(filepath):
                 ws.cell(row=current_row, column=col).border = top_border
             current_row += len(dist_lines)
 
-    wb.save(filepath.replace(".xlsx", "（压缩）.xlsx"))
+    compressed_filepath = filepath.replace(".xlsx", "（压缩）.xlsx")
+    wb.save(compressed_filepath)
+    return compressed_filepath
 
 
 if __name__ == "__main__":