Python - Not repeating code with if statements, is there another way?

Question

Looking to improve my python and coding skills. I have a function that adds a particular timeframe to a time. I pass in:

1M, 7D, 6M, 2H, M etc .. and return the value. I feel like I am repeating myself. Is there a more pythonic approach to this?

def add_timeframe(time, timeframe):                                                                                                                                     

    if 'H' in timeframe:                                                                                                                                                
    ┆   try:                                                                                                                                                            
    ┆   ┆   period = int(re.sub('\D', '', timeframe))                                                                                                                   
    ┆   ┆   return convert_datetime(time + datetime.timedelta(hours=period))                                                                                            

    ┆   except ValueError:                                                                                                                                              
    ┆   ┆   return convert_datetime(time + datetime.timedelta(hours=1))                                                                                                 


    if 'D' in timeframe:                                                                                                                                                
    ┆   try:                                                                                                                                                            
    ┆   ┆   period = int(re.sub('\D', '', timeframe))                                                                                                                   
    ┆   ┆   return convert_datetime(time + datetime.timedelta(days=period))                                                                                             

    ┆   except ValueError:                                                                                                                                              
    ┆   ┆   return convert_datetime(time + datetime.timedelta(days=1))                                                                                                  


    if 'W' in timeframe:                                                                                                                                                
    ┆   try:                                                                                                                                                            
    ┆   ┆   period = int(re.sub('\D', '', timeframe))                                                                                                                   
    ┆   ┆   return convert_datetime(time + datetime.timedelta(weeks=period))                                                                                            

    ┆   except ValueError:                                                                                                                                              
    ┆   ┆   return convert_datetime(time + datetime.timedelta(weeks=period))                                                                                            


    if 'M' in timeframe:                                                                                                                                                
    ┆   try:                                                                                                                                                            
    ┆   ┆   period = int(re.sub('\D', '', timeframe))                                                                                                                   
    ┆   ┆   return convert_datetime(time + datetime.timedelta(days=365/12*period))                                                                                      

    ┆   except ValueError:                                                                                                                                              
    ┆   ┆   return convert_datetime(time + datetime.timedelta(days=365/12))

Answer 1

I usually avoid lots of ifs by using dictionaries. I map each condition to a dictionary and execute. Here is my first take:

I created a function for adding months as timedelta does not have it. I then use re to get digit and letter as tuples. So '4M' would be ('4','M'). Then I would map M to month addition, so 4 *(add months function), W for add weeks etc

import calendar
import datetime
import re


# add month hack
def add_month(num_months, date=None):
    '''Add N months'''

    assert num_months > 0, 'Positive N only'

    if date is None:
        date = datetime.datetime.now()

    for num in range(num_months):
        month_days = calendar.monthrange(date.year, date.month)[1]
        dt = date + datetime.timedelta(days=month_days)

        if dt.day != date.day:
             dt.replace(day=1) - datetime.timedelta(days=1)
        else:
            dt

        date = dt
    return dt



def delta(data, pattern):
    # dict instead of lots of ifs

    time_convert = {'M': lambda x : add_month(x),
                   'W': lambda x :datetime.timedelta(weeks=x),
                   'D': lambda x: datetime.timedelta(days=x),
                   'H': lambda x: datetime.timedelta(hours=x),
                          }


    _ = [re.match(pattern, item).groups() for item in data]

    return [time_convert.get(letter)(int(number))for number, letter in _] 

# test 1
data = ['1M', '7D', '4M', '2H']
pattern = '(\d+)(\w+)'

s = delta(data, pattern)
print(s)

If we are expecting unclean data, we need to create a data preparation function that will ensure that our data is in a format we want digitsletter(s). As our code will fail if we receive just letter data. Instead of try-catch, If we have just a letter, we could add a 1. This would be my take 2:

import calendar
import datetime
import re


# add month hack
def add_month(num_months, date=None):
    '''Add N months'''

    assert num_months > 0, 'Positive N only'

    if date is None:
        date = datetime.datetime.now()

    for num in range(num_months):
        month_days = calendar.monthrange(date.year, date.month)[1]
        dt = date + datetime.timedelta(days=month_days)

        if dt.day != date.day:
             dt.replace(day=1) - datetime.timedelta(days=1)
        else:
            dt

        date = dt
    return dt


def data_prep(data, check_patter='\d+'):
    '''Our data preparation happens here'''

    _ = [bool(re.search(check_patter,item)) for item in data]

    for index, truth in enumerate(_):
        if not truth:
            data[index] = '1'+data[index]

    return data




def delta(data, pattern):

    time_convert = {'M': lambda x : add_month(x),
                'W': lambda x :datetime.timedelta(weeks=x),
                'D': lambda x: datetime.timedelta(days=x),
                'H': lambda x: datetime.timedelta(hours=x),
                       }


    # clean data. if M --> 1M
    data = data_prep(data)
    _ = [re.match(pattern, item).groups() for item in data]

    return [time_convert.get(letter)(int(number))for number, letter in _] 


data = ['1M', '7D', '4M', '2H','H']

pattern = '(\d+)(\w+)'


s = delta(data, pattern)

print(s)

It is in data_prep function where you will deal with all possible uncleanliness:)

Answer 2

You can use re to extract both the number portion and all the period-chars in a single go..

>>> import re
>>> inp = ["1M", "7D", "6M", "2H", "M"]
>>> [re.findall('(\d)?(M|D|H)', x) for x in inp]
[[('1', 'M')], [('7', 'D')], [('6', 'M')], [('2', 'H')], [('', 'M')]]

>>> extracted = [re.findall('(\d)?(M|D|H)', x) for x in inp]
>>> [(int(x[0][0] or '1'), x[0][1]) for x in extracted if x] # Filter out invalids.
[(1, 'M'), (7, 'D'), (6, 'M'), (2, 'H'), (1, 'M')]

You can then use convert_datetime(..) and other stuff you are doing in your original code.

PS: I'd perform more error-checking -- the code above is to only suggest a slightly more pythonic way of doing the same thing.