Commits

Robert Brewer committed 3b97f7c

Oops. I needed those two classes after all

  • Participants
  • Parent commits 2aa5ab5

Comments (0)

Files changed (1)

File flowrate/flows.py

 import flowrate
 
 
+class month(object):
+    """A given month in a given year. Like datetime.date without a .day field."""
+
+    __slots__ = ('year', 'month')
+
+    def __init__(self, year, month):
+        self.year = year
+        self.month = month
+
+    @classmethod
+    def fromdate(cls, dt):
+        return cls(dt.year, dt.month)
+
+    def todate(self, day=1):
+        return datetime.date(self.year, self.month, day)
+
+    def __add__(self, other):
+        # Add an integer number of months to self.
+        y, m = divmod(self.month + other, 12)
+        return month(self.year + y, m)
+
+    def __sub__(self, other):
+        if isinstance(other, month):
+            # Subtract a month() instance to return an integer diff of months
+            y, m = divmod(self.month - other.month, 12)
+            return (y * 12) + m
+        else:
+            # Subtract an integer number of months from self.
+            y, m = divmod(self.month - other, 12)
+            return month(self.year + y, m)
+
+    def __cmp__(self, other):
+        return cmp((self.year, self.month), (other.year, other.month))
+
+
+class week(object):
+    """A given week in a given year. Like datetime.date with a .week field."""
+
+    __slots__ = ('year', 'week')
+
+    def __init__(self, year, week):
+        self.year = year
+        self.week = week
+
+    @classmethod
+    def fromdate(cls, dt):
+        y, w, d =  dt.isocalendar()
+        return cls(y, w)
+
+    def todate(self, day=1):
+        jan1 = datetime.date(self.year, 1, 1)
+        day1ofweek1 = jan1 + datetime.timedelta(days=8 - jan1.isocalendar()[2])
+        return (day1ofweek1 +
+                datetime.timedelta(
+                    days=((self.week - 1) * 7) + (day - 1)
+                ))
+
+    def __add__(self, other):
+        # Add an integer number of weeks to self.
+        f = self.todate() + datetime.timedelta(days=other * 7)
+        y, w, d = f.isocalendar()
+        return week(y, w)
+
+    def __sub__(self, other):
+        if isinstance(other, week):
+            # Subtract a week() instance to return an integer diff of weeks
+            return ((self.todate() - other.todate()).days / 7)
+        else:
+            # Subtract an integer number of weeks from self.
+            f = self.todate() - datetime.timedelta(days=other * 7)
+            y, w, d = f.isocalendar()
+            return week(y, w)
+
+    def __cmp__(self, other):
+        return cmp((self.year, self.week), (other.year, other.week))
+
+
 def obligations(flow):
     """Yield obligation rows for the given flow row."""
     # Add dummy transactions for this flow based on its unit