I have a dataframe whose columns are populated using faker, this is the code

import pandas as pd
import numpy as np
import faker

fake = Faker()

def make_workers():
    status_list = ['in', 'out']
    room_list = ['FL1_RM1','FL1_RM2','FL1_RM3', 'FL1_RM4', 'FL2_RM1','FL2_RM2','FL2_RM3','FL2_RM4', 'FL3_RM1', 'FL3_RM2','FL3_RM3', 'FL3_RM4', 'FL4_RM1','FL4_RM2','FL4_RM3', 'FL4_RM4']
    Property = ['B1','B2','B3','B4']
    d1 = datetime.strptime(f'03/01/2022', '%m/%d/%Y')
    d2 = datetime.strptime(f'08/08/2022', '%m/%d/%Y')
    timestamps = pd.date_range(d1, d2, freq="1min")
    fake_workers = [{'ID':fake.random_number(digits=10),
                  'Property num': np.random.choice(Property, p=[0.25, 0.25, 0.25, 0.25]),
                  'Temp':np.random.randint(low=35, high=50), 
                  'noted Date':timestamps[x],
                  'Status':np.random.choice(status_list),
                  'Humidity': np.random.uniform(low=-35.09, high=70.00),
                   
                  'Dust':np.random.randint(low=2, high=5),
                  'CO2 level':np.random.uniform(low= 350.09, high=450.00),
                  'room_id':np.random.choice(room_list)} for x in range(len(timestamps))]
        
    return fake_workers

worker_df = pd.DataFrame(make_workers())
worker_df.head(30)

I want to add a new column called Floor_Number which takes the values from third letter from room_id column, for example if room_id -> FL1_RM2 then Floor_number -> is 1 and so on

whats the way for this

CodePudding user response：

One way is to nest the list comprehension you already have going:

from datetime import datetime

import numpy as np
import pandas as pd
from faker import Faker

fake = Faker()


def make_workers() -> list:
    status_list = ['in', 'out']
    room_list = ['FL1_RM1', 'FL1_RM2', 'FL1_RM3', 'FL1_RM4', 'FL2_RM1', 'FL2_RM2', 'FL2_RM3', 'FL2_RM4', 'FL3_RM1',
                 'FL3_RM2', 'FL3_RM3', 'FL3_RM4', 'FL4_RM1', 'FL4_RM2', 'FL4_RM3', 'FL4_RM4']
    Property = ['B1', 'B2', 'B3', 'B4']
    d1 = datetime.strptime('03/01/2022', '%m/%d/%Y')
    d2 = datetime.strptime('08/08/2022', '%m/%d/%Y')
    timestamps = pd.date_range(d1, d2, freq="1min")
    return [{**elem, **{"Floor_Number": elem.get("room_id")[2]}} for elem in [
        {'ID'          : fake.random_number(digits=10),
         'Property num': np.random.choice(Property, p=[0.25, 0.25, 0.25, 0.25]),
         'Temp'        : np.random.randint(low=35, high=50),
         'noted Date'  : timestamps[x],
         'Status'      : np.random.choice(status_list),
         'Humidity'    : np.random.uniform(low=-35.09, high=70.00),
         'Dust'        : np.random.randint(low=2, high=5),
         'CO2 level'   : np.random.uniform(low=350.09, high=450.00),
         'room_id'     : np.random.choice(room_list)} for x in range(len(timestamps))]]


worker_df = pd.DataFrame(make_workers())
print(worker_df.head(30))

Readability could be improved here, perhaps by creating the entire dataframe within make_workers() or splitting this process/function into parts.

Output