Skip to content

Commit

Permalink
fix: Fixed list index out of range error in lookup_area function.
Browse files Browse the repository at this point in the history
  • Loading branch information
GabrielKuka committed Oct 7, 2022
1 parent 7c745d9 commit 1966e4c
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 10 deletions.
24 changes: 15 additions & 9 deletions entsoe/mappings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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',
Expand Down
2 changes: 1 addition & 1 deletion test.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,4 +41,4 @@
lst.append(s)

result = pd.concat(lst)
result.to_csv('result.csv')
result.to_csv('result.csv')

0 comments on commit 1966e4c

Please sign in to comment.