Update src/extractor.py

src/extractor.py

@@ -4,9 +4,7 @@ GST_REGEX = r"\d{2}[A-Z]{5}\d{4}[A-Z]{1}[A-Z\d]{1}[Z]{1}[A-Z\d]{1}"
 
 def extract_amount(line):
     nums = re.findall(r"\d+\.\d+|\d+", line.replace(",", ""))
-    if nums:
-        return float(nums[-1])
-    return 0.0
+    return float(nums[-1]) if nums else 0.0
 
 def extract_fields(text):
     lines = text.split("\n")
@@ -28,36 +26,32 @@ def extract_fields(text):
         data["GSTIN"] = gst_match.group()
 
     for line in lines:
-        lower = line.lower()
+        l = line.lower()
 
-        if "invoice" in lower and not data["Invoice No"]:
-            words = line.split()
-            data["Invoice No"] = words[-1]
+        if "invoice" in l and not data["Invoice No"]:
+            data["Invoice No"] = line.split()[-1]
 
-        if "date" in lower and not data["Date"]:
-            words = line.split()
-            data["Date"] = words[-1]
+        if "date" in l and not data["Date"]:
+            data["Date"] = line.split()[-1]
 
-        if "cgst" in lower:
+        if "cgst" in l:
             data["CGST"] = extract_amount(line)
 
-        if "sgst" in lower:
+        if "sgst" in l:
             data["SGST"] = extract_amount(line)
 
-        if "igst" in lower:
+        if "igst" in l:
             data["IGST"] = extract_amount(line)
 
-        if "taxable" in lower:
+        if "taxable" in l:
             data["Taxable"] = extract_amount(line)
 
-        if "total" in lower:
+        if "total" in l:
             data["Total"] = extract_amount(line)
 
-    # Vendor fallback (first non-empty line)
-    if not data["Vendor"]:
-        for l in lines:
-            if len(l.strip()) > 5:
-                data["Vendor"] = l.strip()
-                break
+    for line in lines:
+        if len(line.strip()) > 5:
+            data["Vendor"] = line.strip()
+            break
 
     return data