I have this function that returns a dataframe of values given a set of numbers:
def metadata(*numbers):
"""Define the state FIPS code and state name from a given state abbreviation.
:param number num -> EIA 930 BA number
:return: color: str -> BA short name
:return: shape: str -> BA long name
"""
results = []
color = ""
shape = ""
for number in numbers:
if number == 1: color = 'green'; shape = 'square'
elif number == 9: color = 'white'; shape = 'circle'
elif number == 7: color = 'red'; shape = 'hexagon'
elif number == 13: color = 'purple'; shape = 'triangle'
elif number == 3: color = 'blue'; shape = 'pentagon'
else: color = None; shape = None
results.append([number, color, shape])
df = pd.DataFrame(results, columns=['number', 'color', 'shape'])
return df
test = metadata(1, 2, 9, 13)
print(test)
This produces the desired result:
number color shape
0 1 green square
1 2 None None
2 9 white circle
3 13 purple triangle
However when I try to input a list of numbers from a dataframe I get an incorrect result:
df_sample = pd.DataFrame(data=np.array([[1, 'Tim'], [2, 'Jackie'], [9, 'Sam'],[13, 'Greg']])
, columns=['id', 'name'])
number = df_sample['id'].tolist()
test_2 = metadata(number)
print(test_2)
Resulting in:
number color shape
0 [1, 2, 9, 13] None None
I know the issue lies within extracting the id column from the df_sample however I am not sure how I can extract the data so it is (1, 2, 9, 13) instead of a list [1, 2, 9, 13]. I have tried removing the .tolist but it just gives a subset of the dataframe including the index which I don't want. Any ideas on how to get the function to read in a dataframe column?