test

test_atc_2.py

import pandas as pd
from io import StringIO

import datetime 


def to_df(table_data):

    df = pd.read_csv(table, delimiter=",")
    return df

def date_parse(d):
    d = str(d)
    d, m, y = d[:2], d[2:4], d[4:]
    print(d, m, y)
    # convert to datetime
    return datetime.datetime(int(y), int(m), int(d))


def process(table_data):
    df = to_df(table_data)
    df['RentedDate'] = df['RentedDate'].apply(date_parse)
    df['Duration'] = df['Duration'].fillna(0)
    df['Duration'] = df['Duration'].apply(lambda x: datetime.timedelta(x))

    return df


def rent_filter(df, csv_name="high.csv"):
    car_rent_count = dict(df[df['Status'] == 'Occupied'].value_counts("Car"))
    low = []
    high = []
    for car in car_rent_count.keys():
        if car_rent_count[car] <= 2:
            low.append(car)
        else:
            high.append(car)

    df[df['Car'].isin(low)].to_csv('low.csv')
    df[df['Car'].isin(high)].to_csv(csv_name)
    return df


def get_returned_date(df):
    df = df.sort_values("RentedDate")
    df['ReturnDate'] = df['RentedDate'] + df['Duration']
    return df.to_json()


if __name__ == '__main__':
    # usually already sit in database, so I will need to write SQL instead
    table = StringIO("""
    Car,Status,RentedDate,Duration
    Car1,Available,11022012,
    Car2,Occupied,15052019,2
    Car3,Occupied,16052020,1
    Car1,Occupied,18062017,3
    Car2,Occupied,20122020,3
    Car1,Available,22042021,
    Car3,Occupied,30042018,1
    Car3,Available,26062019,
    Car1,Occupied,23042021,2
    Car3,Occupied,26062019,2
    Car2,Occupied,30122020,3
    """)
    processed_df = process(table)
    # first problem

    rent_filter(processed_df)
    get_returned_date(processed_df)