""" Management command: export_referral_codes_to_excel Reads the legacy Excel workbook, looks up each user by mobile (col F), and writes their system referral_code into col D. Rows where the user is not found in the system are left unchanged. Usage: python manage.py export_referral_codes_to_excel \ --file "C:/Users/DN761673/Desktop/TWM_Legacy register_data.xlsx" # Preview without saving: python manage.py export_referral_codes_to_excel \ --file "C:/Users/DN761673/Desktop/TWM_Legacy register_data.xlsx" \ --dry-run # Write to a separate file instead of overwriting: python manage.py export_referral_codes_to_excel \ --file "C:/Users/DN761673/Desktop/TWM_Legacy register_data.xlsx" \ --output "C:/Users/DN761673/Desktop/TWM_Legacy_with_codes.xlsx" """ from __future__ import annotations import html import re import zipfile from pathlib import Path from xml.etree import ElementTree as ET from django.core.management.base import BaseCommand, CommandError from apps.accounts.models import User class Command(BaseCommand): help = ( "Look up each legacy-workbook user's referral_code from the system " "and write it into col D of the workbook." ) # ------------------------------------------------------------------ # CLI arguments # ------------------------------------------------------------------ def add_arguments(self, parser): parser.add_argument("--file", required=True, help="Absolute path to legacy workbook (.xlsx)") parser.add_argument("--sheet", default=None, help="Optional sheet name. Defaults to active sheet.") parser.add_argument("--output", default=None, help="Output path. Defaults to overwriting --file.") parser.add_argument("--dry-run", action="store_true", default=False, help="Print changes without saving the file.") # ------------------------------------------------------------------ # Helpers: ZIP/XML resilient reader # ------------------------------------------------------------------ @staticmethod def _read_zip_text(zf: zipfile.ZipFile, name: str) -> str: original_crc = zipfile.ZipExtFile._update_crc zipfile.ZipExtFile._update_crc = lambda self, newdata: None try: return zf.read(name).decode("utf-8", errors="ignore") finally: zipfile.ZipExtFile._update_crc = original_crc @staticmethod def _attr_value(attrs: str, name: str) -> str: match = re.search(rf'{name}="([^"]*)"', attrs) return match.group(1) if match else "" def _resolve_sheet_path(self, zf: zipfile.ZipFile, sheet_name: str | None) -> str: workbook_xml = self._read_zip_text(zf, "xl/workbook.xml") workbook_rels = self._read_zip_text(zf, "xl/_rels/workbook.xml.rels") ns_main = "http://schemas.openxmlformats.org/spreadsheetml/2006/main" ns_r = "http://schemas.openxmlformats.org/officeDocument/2006/relationships" ns_rel = "http://schemas.openxmlformats.org/package/2006/relationships" wb_root = ET.fromstring(workbook_xml) rels_root = ET.fromstring(workbook_rels) first_rid = None target_rid = None for sheet in wb_root.findall(f"{{{ns_main}}}sheets/{{{ns_main}}}sheet"): rid = sheet.attrib.get(f"{{{ns_r}}}id") if first_rid is None: first_rid = rid if sheet_name and sheet.attrib.get("name") == sheet_name: target_rid = rid break rid = target_rid or first_rid if not rid: raise CommandError("Could not resolve a worksheet in the workbook.") for rel in rels_root.findall(f"{{{ns_rel}}}Relationship"): if rel.attrib.get("Id") == rid: target = rel.attrib.get("Target", "") if target.startswith("/"): target = target.lstrip("/") if not target.startswith("xl/"): target = f"xl/{target}" return target raise CommandError("Could not find worksheet relationship target.") def _load_shared_strings(self, zf: zipfile.ZipFile) -> list[str]: if "xl/sharedStrings.xml" not in zf.namelist(): return [] xml = self._read_zip_text(zf, "xl/sharedStrings.xml") try: root = ET.fromstring(xml) except ET.ParseError: return [] ns = "http://schemas.openxmlformats.org/spreadsheetml/2006/main" result: list[str] = [] for si in root.findall(f"{{{ns}}}si"): texts = [t.text or "" for t in si.findall(f".//{{{ns}}}t")] result.append("".join(texts)) return result def _extract_cell_value(self, cell_xml: str, shared: list[str]) -> tuple[str, str]: attrs_match = re.match(r"]*)>", cell_xml) attrs = attrs_match.group(1) if attrs_match else "" ref = self._attr_value(attrs, "r") cell_type = self._attr_value(attrs, "t") if cell_type == "inlineStr": m = re.search(r"]*>(.*?)", cell_xml, re.S) return ref, html.unescape(m.group(1)) if m else "" if cell_type == "s": m = re.search(r"(.*?)", cell_xml, re.S) if not m: return ref, "" idx_text = html.unescape(m.group(1)).strip() if idx_text.isdigit(): idx = int(idx_text) if 0 <= idx < len(shared): return ref, shared[idx] return ref, idx_text m = re.search(r"(.*?)", cell_xml, re.S) return ref, html.unescape(m.group(1)) if m else "" def _iter_rows(self, sheet_xml: str, shared: list[str]): for row_m in re.finditer(r']*r="(\d+)"[^>]*>(.*?)', sheet_xml, re.S): row_num = int(row_m.group(1)) row_xml = row_m.group(2) values: dict[str, str] = {} for cell_m in re.finditer(r']*>.*?', row_xml, re.S): ref, val = self._extract_cell_value(cell_m.group(0), shared) if ref: col = re.match(r"([A-Z]+)", ref) values[col.group(1) if col else ref] = val yield row_num, values # ------------------------------------------------------------------ # Helpers: column index conversion # ------------------------------------------------------------------ @staticmethod def _col_to_index(col: str) -> int: col = col.upper() result = 0 for ch in col: result = result * 26 + (ord(ch) - ord('A') + 1) return result @staticmethod def _index_to_col(index: int) -> str: out = "" while index > 0: index, rem = divmod(index - 1, 26) out = chr(ord('A') + rem) + out return out # ------------------------------------------------------------------ # Helpers: phone normalisation # ------------------------------------------------------------------ @staticmethod def _clean_phone(value) -> str: raw = str(value).strip() if value else "" if raw.endswith(".0") and raw.replace(".", "", 1).isdigit(): raw = raw[:-2] digits = "".join(ch for ch in raw if ch.isdigit()) return digits if digits else raw # ------------------------------------------------------------------ # Writer (openpyxl) # ------------------------------------------------------------------ def _save_workbook(self, rows: dict[int, dict[str, str]], output_path: Path): try: from openpyxl import Workbook except ImportError: raise CommandError("openpyxl is required to write the output workbook. pip install openpyxl") # Determine max column max_col = 1 for vals in rows.values(): for col in vals: max_col = max(max_col, self._col_to_index(col)) wb = Workbook() ws = wb.active for row_num in sorted(rows.keys()): vals = rows[row_num] for col_idx in range(1, max_col + 1): ws.cell(row=row_num, column=col_idx).value = vals.get( self._index_to_col(col_idx), "" ) output_path.parent.mkdir(parents=True, exist_ok=True) wb.save(str(output_path)) # ------------------------------------------------------------------ # Main handler # ------------------------------------------------------------------ def handle(self, *args, **options): file_path = Path(options["file"]) sheet_name = options.get("sheet") dry_run: bool = bool(options.get("dry_run")) output_raw = options.get("output") output_path = Path(output_raw) if output_raw else file_path if not file_path.exists(): raise CommandError(f"File not found: {file_path}") if file_path.suffix.lower() != ".xlsx": raise CommandError("Only .xlsx files are supported.") # ── Parse workbook ────────────────────────────────────────── with zipfile.ZipFile(file_path) as zf: sheet_path = self._resolve_sheet_path(zf, sheet_name) sheet_xml = self._read_zip_text(zf, sheet_path) shared = self._load_shared_strings(zf) complete_rows = list(self._iter_rows(sheet_xml, shared)) if not complete_rows: raise CommandError("No readable rows found in the workbook sheet.") rows_by_number: dict[int, dict[str, str]] = { rn: dict(rv) for rn, rv in complete_rows } # ── Pre-load all users (mobile → referral_code) ────────────── self.stdout.write("Loading users from database...") mobile_to_code: dict[str, str] = {} for user in User.objects.exclude(mobile="").exclude(mobile__isnull=True).only("mobile", "referral_code"): if user.mobile and user.referral_code: mobile_to_code[user.mobile.strip()] = user.referral_code.strip() self.stdout.write(f"Loaded {len(mobile_to_code)} users with referral codes.") # ── Match rows and update col D ─────────────────────────────── updated = 0 not_found = 0 already_correct = 0 no_mobile = 0 for row_num in sorted(rows_by_number.keys()): if row_num == 1: continue # header row row = rows_by_number[row_num] mobile = self._clean_phone(row.get("F", "")) if not mobile: no_mobile += 1 continue code = mobile_to_code.get(mobile) if code is None: not_found += 1 continue existing_code = row.get("D", "").strip() if existing_code == code: already_correct += 1 continue if dry_run: self.stdout.write( f" Row {row_num}: mobile={mobile} D: '{existing_code}' -> '{code}'" ) else: rows_by_number[row_num]["D"] = code updated += 1 # ── Save ────────────────────────────────────────────────────── if not dry_run and updated > 0: self._save_workbook(rows_by_number, output_path) self.stdout.write(self.style.SUCCESS( f"Saved workbook to: {output_path}" )) elif not dry_run: self.stdout.write("No changes needed — workbook not modified.") # ── Summary ─────────────────────────────────────────────────── self.stdout.write("") self.stdout.write(self.style.SUCCESS("Done.")) self.stdout.write(f" Updated (col D written) : {updated}") self.stdout.write(f" Already correct : {already_correct}") self.stdout.write(f" User not found in system : {not_found}") self.stdout.write(f" Rows with no mobile : {no_mobile}") if dry_run: self.stdout.write(self.style.WARNING(" DRY RUN — file was NOT saved."))