0

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?

KC Ray
  • 67
  • 8
  • 2
    Your functions does not take a list... You should unpack into args `test_2 = metadata(*number)` or update the function definition `def metadata(numbers)` – Henry Ecker Nov 03 '21 at 21:57

0 Answers0