Update splitter.py

splitter.py

@@ -50,7 +50,7 @@ html_subject = """
                 -webkit-text-fill-color: transparent;
                 margin: 0;
                 text-shadow: 0 2px 5px rgba(0, 0, 0, 0.4);
-            ">Upload a CSV file</h3>
+            ">Upload a CSV or XLSX file</h3>
         </div>
     </body>
     </html>
@@ -58,47 +58,56 @@ html_subject = """
 
 st.markdown(html_subject, unsafe_allow_html=True)
 
-# Step 1: Upload CSV file
+# Step 1: Upload CSV or XLSX file
 uploaded_file = st.file_uploader("", type=["csv", "xlsx"])
 
 if uploaded_file is not None:
     try:
-        # Try reading the CSV file with utf-8 encoding first
-        df = pd.read_csv(uploaded_file, encoding='utf-8')
-    except UnicodeDecodeError:
-        # Fallback to ISO-8859-1 encoding if utf-8 fails
-        df = pd.read_csv(uploaded_file, encoding='ISO-8859-1')
+        # Check file type and read accordingly
+        if uploaded_file.name.endswith('.csv'):
+            # Try reading the CSV file with utf-8 encoding first
+            df = pd.read_csv(uploaded_file, encoding='utf-8')
+        elif uploaded_file.name.endswith('.xlsx'):
+            # Read the Excel file
+            df = pd.read_excel(uploaded_file, engine='openpyxl')
+        else:
+            st.error("Unsupported file format. Please upload a CSV or XLSX file.")
+            df = None
+    except Exception as e:
+        st.error(f"An error occurred while reading the file: {e}")
+        df = None
 
-    if 'Invoice_No' not in df.columns:
-        st.error("The uploaded CSV does not contain the 'Invoice_No' column.")
-    else:
-        # Step 3: Split the DataFrame based on unique Invoice_No values
-        invoice_groups = df.groupby('Invoice_No')
-        
-        # Step 4: Generate CSV files for each Invoice_No and add them to a ZIP archive
-        zip_buffer_csv = BytesIO()
-        with ZipFile(zip_buffer_csv, "a") as zip_file:
-            for invoice_no, group in invoice_groups:
-                csv_buffer = BytesIO()
-                group.to_csv(csv_buffer, index=False)
-                zip_file.writestr(f"{invoice_no}.csv", csv_buffer.getvalue())
+    if df is not None:
+        if 'Invoice_No' not in df.columns:
+            st.error("The uploaded file does not contain the 'Invoice_No' column.")
+        else:
+            # Step 3: Split the DataFrame based on unique Invoice_No values
+            invoice_groups = df.groupby('Invoice_No')
+            
+            # Step 4: Generate CSV files for each Invoice_No and add them to a ZIP archive
+            zip_buffer_csv = BytesIO()
+            with ZipFile(zip_buffer_csv, "a") as zip_file:
+                for invoice_no, group in invoice_groups:
+                    csv_buffer = BytesIO()
+                    group.to_csv(csv_buffer, index=False)
+                    zip_file.writestr(f"{invoice_no}.csv", csv_buffer.getvalue())
 
-        zip_buffer_csv.seek(0)  # Move to the beginning of the buffer
-        
-        # Step 5: Generate Excel files for each Invoice_No and add them to a ZIP archive
-        zip_buffer_xlsx = BytesIO()
-        with ZipFile(zip_buffer_xlsx, "a") as zip_file:
-            for invoice_no, group in invoice_groups:
-                excel_buffer = BytesIO()
-                with pd.ExcelWriter(excel_buffer, engine='xlsxwriter') as writer:
-                    group.to_excel(writer, index=False, sheet_name=str(invoice_no))
-                excel_buffer.seek(0)
-                zip_file.writestr(f"{invoice_no}.xlsx", excel_buffer.getvalue())
+            zip_buffer_csv.seek(0)  # Move to the beginning of the buffer
+            
+            # Step 5: Generate Excel files for each Invoice_No and add them to a ZIP archive
+            zip_buffer_xlsx = BytesIO()
+            with ZipFile(zip_buffer_xlsx, "a") as zip_file:
+                for invoice_no, group in invoice_groups:
+                    excel_buffer = BytesIO()
+                    with pd.ExcelWriter(excel_buffer, engine='xlsxwriter') as writer:
+                        group.to_excel(writer, index=False, sheet_name=str(invoice_no))
+                    excel_buffer.seek(0)
+                    zip_file.writestr(f"{invoice_no}.xlsx", excel_buffer.getvalue())
 
-        zip_buffer_xlsx.seek(0)  # Move to the beginning of the buffer
+            zip_buffer_xlsx.seek(0)  # Move to the beginning of the buffer
 
-        # Step 6: Provide download buttons for both ZIP files (CSV and Excel)
-        button_styles = """
+            # Step 6: Provide download buttons for both ZIP files (CSV and Excel)
+            button_styles = """
                 <style>
                 div.stDownloadButton > button {
                     color: #ffffff; /* Text color */
@@ -119,18 +128,19 @@ if uploaded_file is not None:
                 }
                 </style>
             """
-        st.markdown(button_styles, unsafe_allow_html=True)
-        # Download buttons
-        st.download_button(
-            label="Download Invoice ZIP (CSV)",
-            data=zip_buffer_csv,
-            file_name="invoices_csv.zip",
-            mime="application/zip"
-        )
+            st.markdown(button_styles, unsafe_allow_html=True)
+            
+            # Download buttons
+            st.download_button(
+                label="Download Invoice ZIP (CSV)",
+                data=zip_buffer_csv,
+                file_name="invoices_csv.zip",
+                mime="application/zip"
+            )
 
-        st.download_button(
-            label="Download Invoice ZIP (Excel)",
-            data=zip_buffer_xlsx,
-            file_name="invoices_xlsx.zip",
-            mime="application/zip"
-        )
+            st.download_button(
+                label="Download Invoice ZIP (Excel)",
+                data=zip_buffer_xlsx,
+                file_name="invoices_xlsx.zip",
+                mime="application/zip"
+            )