refactor: Melhorias nas funções de validação de CNPJ

brutils/cnpj.py

@@ -1,7 +1,10 @@
+import random
+import re
 from itertools import chain
-from random import choices, randint
 from string import ascii_uppercase, digits
 
+CNPJ_RE = re.compile(r"^[A-Z0-9]{12}[0-9]{2}")
+
 # FORMATTING
 ############
 
@@ -31,7 +34,7 @@ def sieve(dirty: str) -> str:
        backward compatibility.
     """
 
-    return "".join(filter(lambda char: char not in "./-", dirty))
+    return re.sub(r"[\.\/\-]", "", dirty)
 
 
 def remove_symbols(dirty: str) -> str:
@@ -84,13 +87,9 @@ def display(cnpj: str) -> str | None:
        backward compatibility.
     """
 
-    if (
-        len(cnpj) != 14
-        or not _is_alphanumeric(cnpj[:12])
-        or not cnpj[12:].isdigit()
-        or len(set(cnpj)) == 1
-    ):
+    if not re.search(CNPJ_RE, cnpj):
         return None
+
     return "{}.{}.{}/{}-{}".format(
         cnpj[:2], cnpj[2:5], cnpj[5:8], cnpj[8:12], cnpj[12:]
     )
@@ -130,27 +129,6 @@ def format_cnpj(cnpj: str) -> str | None:
 ############
 
 
-def _is_alphanumeric(cnpj: str) -> bool:
-    """
-    Checks whether all characters are digits or uppercase letters.
-
-    Args:
-        cnpj (str): The CNPJ string to be validated.
-
-    Returns:
-        bool: True if all characters are either digits or uppercase letters,
-        False otherwise.
-
-    Example:
-        >>> _is_alphanumeric("035ABC1400Z142")
-        True
-        >>> _is_alphanumeric("0011-22200013!")
-        False
-    """
-
-    return all(char in (digits + ascii_uppercase) for char in cnpj)
-
-
 def validate(cnpj: str) -> bool:
     """
     Validates a CNPJ (Brazilian Company Registration Number) by comparing its
@@ -177,14 +155,9 @@ def validate(cnpj: str) -> bool:
        This method should not be used in new code and is only provided for
        backward compatibility.
     """
-
-    if (
-        len(cnpj) != 14
-        or not _is_alphanumeric(cnpj[:12])
-        or not cnpj[12:].isdigit()
-        or len(set(cnpj)) == 1
-    ):
+    if not re.search(CNPJ_RE, cnpj):
         return False
+
     return all(
         _hashdigit(cnpj, i + 13) == int(v) for i, v in enumerate(cnpj[12:])
     )
@@ -239,26 +212,12 @@ def generate(branch: int | str = 1, alphanumeric: bool = False) -> str:
         >>> generate(branch="AB12", alphanumeric=True)
         "NX9K79E2AB1200"
     """
+    final_branch = str(branch)[:4].zfill(4)
 
-    if alphanumeric:
-        branch = str(branch)
-        branch = branch[:4] if len(branch) >= 4 else branch.zfill(4)
-        branch = (
-            "0001"
-            if branch == "0000" or not _is_alphanumeric(branch)
-            else branch
-        )
-        base = "".join(choices(digits * 3 + ascii_uppercase, k=8)) + branch
-
-        return base + _checksum(base)
-
-    branch = int(branch)
-    branch %= 10000
-    branch += int(branch == 0)
-    branch = str(branch).zfill(4)
-    base = str(randint(0, 99999999)).zfill(8) + branch
+    character_choices = digits + ascii_uppercase if alphanumeric else digits
+    cnpj_base = "".join(random.sample(character_choices, 8)) + final_branch
 
-    return base + _checksum(base)
+    return cnpj_base + _checksum(cnpj_base)
 
 
 def _hashdigit(cnpj: str, position: int) -> int:

tests/test_cnpj.py

@@ -1,10 +1,10 @@
+import re
 from unittest import TestCase, main
 from unittest.mock import patch
 
 from brutils.cnpj import (
     _checksum,
     _hashdigit,
-    _is_alphanumeric,
     display,
     format_cnpj,
     generate,
@@ -28,18 +28,12 @@ def test_sieve(self):
 
     def test_display(self):
         self.assertEqual(display("00000000000109"), "00.000.000/0001-09")
+        self.assertEqual(display("00000000000000"), "00.000.000/0000-00")
         self.assertEqual(display("12ABC34501DE35"), "12.ABC.345/01DE-35")
         self.assertIsNone(display("12ABC34501DEAA"))
-        self.assertIsNone(display("00000000000000"))
         self.assertIsNone(display("0000000000000"))
         self.assertIsNone(display("0000000000000a"))
 
-    def test__is_alphanumeric(self):
-        self.assertIs(_is_alphanumeric("12ABC34501DE35"), True)
-        self.assertIs(_is_alphanumeric("12345678910111"), True)
-        self.assertIs(_is_alphanumeric("123456a78b10C1"), False)
-        self.assertIs(_is_alphanumeric("12.ABC.345/01DE-35"), False)
-
     def test_validate(self):
         self.assertIs(validate("34665388000161"), True)
         self.assertIs(validate("12ABC34501DE35"), True)
@@ -85,10 +79,28 @@ def test_generate(self):
             self.assertIsNotNone(display(generate()))
         self.assertIs(validate(generate(branch=1234)), True)
 
+    def test_generate_with_branch(self):
+        test_branches = [
+            (1, "0001"),
+            (2, "0002"),
+            ("A", "000A"),
+            ("ABCD", "ABCD"),
+            (1234, "1234"),
+            (12345, "1234"),
+        ]
+        for branch, final_branch in test_branches:
+            generated_cnpj = generate(branch=branch)
+            self.assertTrue(
+                re.search(
+                    rf"[0-9]{{8}}{final_branch}[0-9]{{2}}", generated_cnpj
+                ),
+                msg=(final_branch, generated_cnpj),
+            )
+
     def test_generate_alphanumeric(self):
         for _ in range(10_000):
             generated = generate(alphanumeric=True)
-            self.assertIs(validate(generated), True)
+            self.assertIs(validate(generated), True, msg=generated)
             self.assertIsNotNone(display(generated))
         self.assertIs(
             validate(generate(branch="1234", alphanumeric=True)), True