From 1966e4cb274a5502d96b68a9b1d046e92d41d76d Mon Sep 17 00:00:00 2001
From: Gabriel <gabrie.kuka@gmail.com>
Date: Fri, 7 Oct 2022 20:21:54 +0200
Subject: [PATCH] fix: Fixed list index out of range error in lookup_area
 function.

---
 entsoe/mappings.py | 24 +++++++++++++++---------
 test.py            |  2 +-
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/entsoe/mappings.py b/entsoe/mappings.py
index 4dff8f6..2569a72 100644
--- a/entsoe/mappings.py
+++ b/entsoe/mappings.py
@@ -5,16 +5,18 @@
 def lookup_area(s: Union['Area', str]) -> 'Area':
     if isinstance(s, Area):
         # If it already is an Area object, we're happy
-        area = s
-    else:  # It is a string
-        try:
-            # If it is a "country code" string, we do a lookup
-            area = Area[s]
-        except KeyError:
-            # It is not, it may be a direct code
-            area = [area for area in Area if area.value == s][0]
-    return area
+        return s
+    if isinstance(s, str):
+        # If it is a "country code" string, we do a lookup
+        if Area.has_code(s.upper()):
+            return Area[s.upper()]
 
+        # If it is a "direct code", we do a lookup
+        for area in Area:
+            if area.value == s:
+                return area
+
+    raise ValueError('Invalid country code.')
 
 class Area(enum.Enum):
     """
@@ -45,6 +47,10 @@ def tz(self):
     def code(self):
         return self.value
 
+    @classmethod
+    def has_code(cls, code:str)->bool:
+        return code in cls.__members__ 
+
     # List taken directly from the API Docs
     DE_50HZ =       '10YDE-VE-------2', '50Hertz CA, DE(50HzT) BZA',                    'Europe/Berlin',
     AL =            '10YAL-KESH-----5', 'Albania, OST BZ / CA / MBA',                   'Europe/Tirane',
diff --git a/test.py b/test.py
index 0f4eb94..e0838cc 100644
--- a/test.py
+++ b/test.py
@@ -41,4 +41,4 @@
         lst.append(s)
 
 result = pd.concat(lst)
-result.to_csv('result.csv')
+result.to_csv('result.csv')
\ No newline at end of file