返回特定時間

下方的 which_date() 函式實現返回某一起始時間後特定一段時間的日期。特定的一段時間有兩種不同的格式：一種是不同的天數，如 "1 day" 或 "30 days"，另一種是不同的星期數，如 "2 weeks" 或 "12 weeks"。

from datetime import datetime, timedelta
def which_date(start_date,time):
    end_date=datetime.strptime(start_date,"%Y/%m/%d")
    n=int(time.split()[0])
    t=time.split()[1]   
    if t in 'days':
        end_date=end_date+timedelta(days=n)
    elif t in 'weeks':
        end_date=end_date+timedelta(weeks=n)
    end_date=end_date.strftime("%Y/%m/%d")
    return end_date

def test():
    assert which_date('2016/02/10','35 days') == '2016/03/16'
    assert which_date('2016/12/21','3 weeks') == '2017/01/11'
    assert which_date('2015/01/17','1 week') == '2015/01/24'
    print("All tests completed.")
    
if __name__ == "__main__":
    test()