Some collected cleanups

.travis.yml

@@ -22,6 +22,22 @@ matrix:
           os: linux
           dist: xenial
           env: TOXENV=py39
+        - python: 3.6
+          os: linux
+          dist: trusty
+          env: TOXENV=mypy
+        - python: 3.7
+          os: linux
+          dist: xenial
+          env: TOXENV=mypy
+        - python: 3.8
+          os: linux
+          dist: xenial
+          env: TOXENV=mypy
+        - python: 3.9
+          os: linux
+          dist: xenial
+          env: TOXENV=mypy
 
 install:
   - pip install tox

setup.cfg

@@ -3,3 +3,9 @@ universal=1
 
 [aliases]
 test=pytest
+
+[mypy]
+ignore_missing_imports = True
+warn_return_any = True
+warn_unreachable = True
+warn_unused_ignores = True

setup.py

@@ -65,7 +65,6 @@ def parse_requirements(requirements, ignore=('setuptools',)):
         "Operating System :: Microsoft :: Windows",
         "Operating System :: POSIX",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.5",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",

tests/test_watson.py

@@ -73,15 +73,21 @@ def test_current_with_empty_given_state(config_dir, mocker):
     assert watson.current == {}
 
 
-def test_current_as_running_frame(watson):
+def test_current_as_running_frame(watson) -> None:
     """
     Ensures frame can be created without a stop date.
     Catches #417: editing task in progress throws an exception
     """
     watson.start('foo', tags=['A', 'B'])
 
     cur = watson.current
-    frame = Frame(cur['start'], None, cur['project'], None, cur['tags'])
+    frame = Frame.make_new(
+        start=cur['start'],
+        stop=None,
+        project=cur['project'],
+        id=None,
+        tags=cur['tags'],
+    )
 
     assert frame.stop is None
     assert frame.project == 'foo'

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = flake8,py35,py36,py37,py38,py39
+envlist = flake8,mypy,py36,py37,py38,py39
 skip_missing_interpreters = True
 
 [testenv]
@@ -15,3 +15,8 @@ usedevelop = True
 [testenv:flake8]
 deps = flake8
 commands = flake8 --show-source watson/ tests/ scripts/
+
+
+[testenv:mypy]
+deps = mypy==0.812
+commands = mypy watson/ tests/ scripts/

watson/cli.py

@@ -5,6 +5,7 @@
 import os
 from dateutil import tz
 from functools import reduce, wraps
+from typing import Optional
 
 import arrow
 import click
@@ -85,29 +86,28 @@ def local_tz_info() -> datetime.tzinfo:
 class DateTimeParamType(click.ParamType):
     name = 'datetime'
 
-    def convert(self, value, param, ctx) -> arrow:
-        if value:
-            date = self._parse_multiformat(value)
-            if date is None:
-                raise click.UsageError(
-                    "Could not match value '{}' to any supported date format"
-                    .format(value)
-                )
-            # When we parse a date, we want to parse it in the timezone
-            # expected by the user, so that midnight is midnight in the local
-            # timezone, or respect the TZ environment variable not in UTC.
-            # Cf issue #16.
-            date = date.replace(tzinfo=local_tz_info())
-            # Add an offset to match the week beginning specified in the
-            # configuration
-            if param.name == "week":
-                week_start = ctx.obj.config.get(
-                    "options", "week_start", "monday")
-                date = apply_weekday_offset(
-                    start_time=date, week_start=week_start)
-            return date
-
-    def _parse_multiformat(self, value) -> arrow:
+    def convert(self, value, param, ctx) -> arrow.Arrow:
+        date = self._parse_multiformat(value)
+        if date is None:
+            raise click.UsageError(
+                "Could not match value '{}' to any supported date format"
+                .format(value)
+            )
+        # When we parse a date, we want to parse it in the timezone
+        # expected by the user, so that midnight is midnight in the local
+        # timezone, or respect the TZ environment variable not in UTC.
+        # Cf issue #16.
+        date = date.replace(tzinfo=local_tz_info())
+        # Add an offset to match the week beginning specified in the
+        # configuration
+        if param.name == "week":
+            week_start = ctx.obj.config.get(
+                "options", "week_start", "monday")
+            date = apply_weekday_offset(
+                start_time=date, week_start=week_start)
+        return date
+
+    def _parse_multiformat(self, value) -> Optional[arrow.Arrow]:
         date = None
         for fmt in (None, 'HH:mm:ss', 'HH:mm'):
             try:
@@ -1256,7 +1256,9 @@ def add(watson, args, from_, to, confirm_new_project, confirm_new_tag):
 @click.argument('id', required=False, autocompletion=get_frames)
 @click.pass_obj
 @catch_watson_error
-def edit(watson, confirm_new_project, confirm_new_tag, id):
+def edit(
+    watson, confirm_new_project: bool, confirm_new_tag: bool, id: Optional[str]
+):
     """
     Edit a frame.
 
@@ -1280,8 +1282,13 @@ def edit(watson, confirm_new_project, confirm_new_tag, id):
         frame = get_frame_from_argument(watson, id)
         id = frame.id
     elif watson.is_started:
-        frame = Frame(watson.current['start'], None, watson.current['project'],
-                      None, watson.current['tags'])
+        frame = Frame.make_new(
+            start=watson.current['start'],
+            stop=None,
+            project=watson.current['project'],
+            id=None,
+            tags=watson.current['tags']
+        )
     elif watson.frames:
         frame = watson.frames[-1]
         id = frame.id

watson/frames.py

@@ -1,19 +1,37 @@
 import uuid
+from typing import Iterable, List, NamedTuple, Optional, Union
 
 import arrow
 
-from collections import namedtuple
-
-HEADERS = ('start', 'stop', 'project', 'id', 'tags', 'updated_at')
-
-
-class Frame(namedtuple('Frame', HEADERS)):
-    def __new__(cls, start, stop, project, id, tags=None, updated_at=None,):
+TimeType = Union[str, arrow.Arrow]
+FrameColumnTypes = Union[
+    arrow.Arrow, Optional[arrow.Arrow], str, Optional[str], List[str]
+]
+
+
+class Frame(NamedTuple):
+    start: arrow.Arrow
+    stop: Optional[arrow.Arrow]
+    project: str
+    id: Optional[str]
+    tags: List[str]
+    updated_at: arrow.Arrow
+
+    @classmethod
+    def make_new(
+        cls,
+        start: TimeType,
+        stop: Optional[TimeType],
+        project: str,
+        id: Optional[str],
+        tags: Optional[List[str]] = None,
+        updated_at: Optional[TimeType] = None,
+    ) -> "Frame":
         try:
             if not isinstance(start, arrow.Arrow):
                 start = arrow.get(start)
 
-            if stop and not isinstance(stop, arrow.Arrow):
+            if stop is not None and not isinstance(stop, arrow.Arrow):
                 stop = arrow.get(stop)
 
             if updated_at is None:
@@ -32,8 +50,13 @@ def __new__(cls, start, stop, project, id, tags=None, updated_at=None,):
         if tags is None:
             tags = []
 
-        return super(Frame, cls).__new__(
-            cls, start, stop, project, id, tags, updated_at
+        return cls(
+            start=start,
+            stop=stop,
+            project=project,
+            id=id,
+            tags=tags,
+            updated_at=updated_at,
         )
 
     def dump(self):
@@ -44,7 +67,7 @@ def dump(self):
         return (start, stop, self.project, self.id, self.tags, updated_at)
 
     @property
-    def day(self):
+    def day(self) -> arrow.Arrow:
         return self.start.floor('day')
 
     def __lt__(self, other):
@@ -60,34 +83,35 @@ def __gte__(self, other):
         return self.start >= other.start
 
 
-class Span(object):
-    def __init__(self, start, stop, timeframe='day'):
+class Span:
+    def __init__(self, start: arrow.Arrow, stop: arrow.Arrow, timeframe='day'):
         self.timeframe = timeframe
         self.start = start.floor(self.timeframe)
         self.stop = stop.ceil(self.timeframe)
 
-    def overlaps(self, frame):
+    def overlaps(self, frame: Frame) -> bool:
         return frame.start <= self.stop and frame.stop >= self.start
 
-    def __contains__(self, frame):
+    def __contains__(self, frame: Frame) -> bool:
         return frame.start >= self.start and frame.stop <= self.stop
 
 
-class Frames(object):
-    def __init__(self, frames=None):
+class Frames:
+    def __init__(self, frames: Optional[List[Frame]] = None):
         if not frames:
             frames = []
 
-        rows = [Frame(*frame) for frame in frames]
+        rows = [Frame.make_new(*frame) for frame in frames]
         self._rows = rows
 
         self.changed = False
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self._rows)
 
     def __getitem__(self, key):
-        if key in HEADERS:
+        attributes = Frame.__annotations__.keys()
+        if key in attributes:
             return tuple(self._get_col(key))
         elif isinstance(key, int):
             return self._rows[key]
@@ -119,31 +143,43 @@ def __delitem__(self, key):
         else:
             del self._rows[self._get_index_by_id(key)]
 
-    def _get_index_by_id(self, id):
+    def _get_index_by_id(self, id: str) -> int:
         try:
             return next(
                 i for i, v in enumerate(self['id']) if v.startswith(id)
             )
         except StopIteration:
             raise KeyError("Frame with id {} not found.".format(id))
 
-    def _get_col(self, col):
-        index = HEADERS.index(col)
+    def _get_col(self, col: str) -> Iterable[FrameColumnTypes]:
         for row in self._rows:
-            yield row[index]
+            yield getattr(row, col)
 
     def add(self, *args, **kwargs):
         self.changed = True
         frame = self.new_frame(*args, **kwargs)
         self._rows.append(frame)
         return frame
 
-    def new_frame(self, project, start, stop, tags=None, id=None,
-                  updated_at=None):
+    def new_frame(
+        self,
+        project: str,
+        start: TimeType,
+        stop: Optional[TimeType],
+        tags: Optional[List[str]] = None,
+        id: Optional[str] = None,
+        updated_at: Optional[TimeType] = None,
+    ) -> Frame:
         if not id:
             id = uuid.uuid4().hex
-        return Frame(start, stop, project, id, tags=tags,
-                     updated_at=updated_at)
+        return Frame.make_new(
+            start=start,
+            stop=stop,
+            project=project,
+            id=id,
+            tags=tags,
+            updated_at=updated_at,
+        )
 
     def dump(self):
         return tuple(frame.dump() for frame in self._rows)
@@ -183,5 +219,5 @@ def filter(
                 stop = span.stop if frame.stop > span.stop else frame.stop
                 yield frame._replace(start=start, stop=stop)
 
-    def span(self, start, stop):
+    def span(self, start: arrow.Arrow, stop: arrow.Arrow) -> Span:
         return Span(start, stop)

watson/utils.py

@@ -186,7 +186,7 @@ def get_start_time_for_period(period):
     return start_time
 
 
-def apply_weekday_offset(start_time, week_start):
+def apply_weekday_offset(start_time: arrow.Arrow, week_start) -> arrow.Arrow:
     """
     Apply the offset required to move the start date `start_time` of a week
     starting on Monday to that of a week starting on `week_start`.

watson/watson.py

@@ -22,7 +22,7 @@ class ConfigurationError(CFGParserError, WatsonError):
     pass
 
 
-class Watson(object):
+class Watson:
     def __init__(self, **kwargs):
         """
         :param frames: If given, should be a list representing the