Implementing a custom counting system in Python

Question

In finance, futures contracts are usually represented by their expiry year and month. So for example, 201212 would be 2012 - December.

Some contracts, for example Corn, are only traded some months [3,5,7,9,12], whereas sometimes, you might only want to trade the [12] contract (despite it trading other months too, so you'd trade 201212, 201312 etc).

I'm currently using the the int format to represent these contracts in my system, and using that as a Pandas index.

The tricky thing is, given a contract, I often need to get the next contract (and to a lesser extent, the previous one).

I've written a generator expression that does something like this:

def contract_generator(self, recent=True, year=None, month=None, traded_only=False):
    if year is None:
        year = datetime.datetime.now().year - 1 if recent == True else self.first_contract

    if traded_only is True:
        months = self.trade_only
    else:
        months = self.months_traded

    months = deque(months)

    if month is not None:
        months.rotate(months.index(month)-1)
    while True:
        for month in months:
            yield {
                'year': year,
                'month': month,
                'contract': str(year)+str("%02d" % (month,)),
                'formatted_contract': self.contract_format(year, month),
                'expiration_date': self.expiration_date(year, month)
                }
        year+=1

def next_contract(self, contract):
    c = contract_to_tuple(contract)
    j = self.contract_generator(year = c[0], month = c[1], traded_only=False)
    return next(j)['contract']

def contract_to_tuple(contract):
    contract = str(contract)
    return (int(contract[0:4]),int(contract[4:6]))

(where months_traded & trade_only are the lists to which I referred to in paragraph 2).

Crucially, its buggy and the above doesn't work quite right. I can fix it, but to be honest, I really dislike this approach. There must be a better way.

Ideas:

• Implement using datetime somehow
• Create some kind of custom object, implement arithmetic operators so I can do 201212 + 1 to get the next contract (but would this really be easy with pandas?)
• I don't think python has it baked in, but perhaps I can define a custom numerical base somehow that maps to this behaviour.

Is there an easy/elegant way to do this that already exists? Or do I really need to make it from scratch?

Edit:

My final result:

def nc(self, contract, months=None):
    d = pd.to_datetime(str(contract), format='%Y%m')
    months_traded = deque(self.months_traded)
    months_traded.rotate(-1)
    output_month = months_traded[self.months_traded.index(d.month)]
    output_year = d.year + 1 * d.month==12
    return str(output_year)+str("%02d" % (output_month,))

Answer 1

This should do it:

def next_contract(contract, last=False):
    d = pd.to_datetime(str(contract), format='%Y%m')
    d += pd.offsets.MonthBegin(12 * (last * -2 + 1))
    return int(d.strftime('%Y%m'))

Demo

next_contract(201212)

201312

next_contract(201212, last=True)

201112

Explanation

def next_contract(contract, last=False):
    # convert contract to datetime with assumed format of yyyymm
    d = pd.to_datetime(str(contract), format='%Y%m')
    # use pandas offsets.  I don't care that it's month begin
    # because I'm ditching the day anyway.
    # (last * -2 + 1) equals -1 when last is True and 1 when last is False
    d += pd.offsets.MonthBegin(12 * (last * -2 + 1))
    return int(d.strftime('%Y%m'))

For what it's worth, here is a stub of a class. Honestly, coding up all the handling for the other months should be left as an exercise for you.

class Corn(object):
    def __init__(self, contract):
        self.contract = contract

    def __add__(self, i):
        d = pd.to_datetime(str(self.contract), format='%Y%m')
        d += pd.offsets.MonthBegin(12 * i)
        self.contract = int(d.strftime('%Y%m'))
        return self

    def __sub__(self, i):
        return self.__add__(-i)

    def get_next(self):
        return self + 1

    def get_last(self):
        return self - 1

    def __repr__(self):
        return str(self.contract)

---

corn = Corn(201212)
print(corn + 1)
print(corn.get_next())

201312
201412