完成下方的 which_date() 函数,并返回某一起始时间后特定一段时间的日期

from datetime import datetime,timedelta
import re
def which_date(start_date,time):
    """
    This function takes as input a string depicting a date in YYYY/mm/dd
    format and a string stating a time period in the form of "X day(s)" or
    "Y week(s)". Output should be a string in form YYYY/mm/dd with the date
    that is X days or Y weeks after the initial date.
    """
    
    # Replace this with your code!
    if time.find('day') > 0:
        day = time.split()
        end_date = datetime.strptime(start_date, "%Y/%m/%d") + timedelta(days=int(day[0]))
    else:
        week = time.split( )
        end_date = datetime.strptime(start_date, "%Y/%m/%d") + timedelta(weeks=int(week[0]))
    return end_date.strftime("%Y/%m/%d")
    
def test():
    """
    Here are a few tests to check if your code is working correctly! This
    function will be run when the Test Run button is selected. Additional
    tests that are not part of this function will also be run when the Submit
    Answer button is selected.
    """
    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()