Expose entity additional codes (#2879)

arkid15r · web-flow · commit 515ba78aa38d · 2025-09-01T01:53:43.000Z
diff --git a/holidays/registry.py b/holidays/registry.py
@@ -13,7 +13,7 @@
 import importlib
 from collections.abc import Iterable
 from threading import RLock
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from holidays.holiday_base import HolidayBase
 
@@ -339,35 +339,46 @@ def get_entity(self) -> Optional[HolidayBase]:
     @staticmethod
     def _get_entity_codes(
         container: RegistryDict,
-        entity_length: Union[int, Iterable[int]],
         include_aliases: bool = True,
+        max_code_length: int = 3,
+        min_code_length: int = 2,
     ) -> Iterable[str]:
-        entity_length = {entity_length} if isinstance(entity_length, int) else set(entity_length)
         for entities in container.values():
-            for entity in entities:
-                if len(entity) in entity_length:
-                    yield entity
-                    # Assuming that the alpha-2 code goes first.
-                    if not include_aliases:
-                        break
+            for code in entities[1:]:
+                if min_code_length <= len(code) <= max_code_length:
+                    yield code
+
+                # Stop after the first matching code if aliases are not requested.
+                # Assuming that the alpha-2 code goes first.
+                if not include_aliases:
+                    break
 
     @staticmethod
     def get_country_codes(include_aliases: bool = True) -> Iterable[str]:
         """Get supported country codes.
 
         :param include_aliases:
-            Whether to include entity aliases (e.g. UK for GB).
+            Whether to include entity aliases (e.g. GBR and UK for GB,
+            UKR for UA, USA for US, etc).
         """
-        return EntityLoader._get_entity_codes(COUNTRIES, 2, include_aliases)
+        return EntityLoader._get_entity_codes(
+            COUNTRIES,
+            include_aliases=include_aliases,
+        )
 
     @staticmethod
     def get_financial_codes(include_aliases: bool = True) -> Iterable[str]:
         """Get supported financial codes.
 
         :param include_aliases:
-            Whether to include entity aliases(e.g. TAR for ECB, XNYS for NYSE).
+            Whether to include entity aliases (e.g. B3 for BVMF,
+            TAR for ECB, NYSE for XNYS, etc).
         """
-        return EntityLoader._get_entity_codes(FINANCIAL, (3, 4), include_aliases)
+        return EntityLoader._get_entity_codes(
+            FINANCIAL,
+            include_aliases=include_aliases,
+            max_code_length=4,
+        )
 
     @staticmethod
     def load(prefix: str, scope: dict) -> None:
diff --git a/tests/test_registry.py b/tests/test_registry.py
@@ -108,6 +108,42 @@ def test_financial_str(self):
             "'holidays.financial.ny_stock_exchange.NYSE' class directly.",
         )
 
+    def test_get_country_codes(self):
+        country_codes = set(registry.EntityLoader.get_country_codes(include_aliases=False))
+        for entity_classes in registry.COUNTRIES.values():
+            self.assertNotIn(entity_classes[0], country_codes)
+            self.assertIn(entity_classes[1], country_codes)
+            for code in entity_classes[2:]:
+                self.assertNotIn(code, country_codes)
+
+    def test_get_country_codes_aliases(self):
+        country_codes = set(registry.EntityLoader.get_country_codes(include_aliases=True))
+        for entity_classes in registry.COUNTRIES.values():
+            self.assertNotIn(entity_classes[0], country_codes)
+            for code in entity_classes[1:]:
+                if code.isupper():
+                    self.assertIn(code, country_codes)
+                else:
+                    self.assertNotIn(code, country_codes)
+
+    def test_get_financial_codes(self):
+        financial_codes = set(registry.EntityLoader.get_financial_codes(include_aliases=False))
+        for entity_classes in registry.FINANCIAL.values():
+            self.assertNotIn(entity_classes[0], financial_codes)
+            self.assertIn(entity_classes[1], financial_codes)
+            for code in entity_classes[2:]:
+                self.assertNotIn(code, financial_codes)
+
+    def test_get_financial_codes_aliases(self):
+        financial_codes = set(registry.EntityLoader.get_financial_codes(include_aliases=True))
+        for entity_classes in registry.FINANCIAL.values():
+            self.assertNotIn(entity_classes[0], financial_codes)
+            for code in entity_classes[1:]:
+                if code.isupper():
+                    self.assertIn(code, financial_codes)
+                else:
+                    self.assertNotIn(code, financial_codes)
+
     def test_inheritance(self):
         def create_instance(parent):
             class SubClass(parent):
