add type hints in data.py

1 files changed, 10 insertions, 23 deletions

diff --git a/fietsboek/data.py b/fietsboek/data.py
index bd4222b..94c6c11 100644
--- a/fietsboek/data.py
+++ b/fietsboek/data.py
@@ -8,6 +8,8 @@ import string
 import shutil
 import uuid
 import logging
+from typing import List, BinaryIO, Optional
+from pathlib import Path
 
 from .util import secure_filename
 
@@ -15,16 +17,14 @@ from .util import secure_filename
 LOGGER = logging.getLogger(__name__)
 
 
-def generate_filename(filename):
+def generate_filename(filename: Optional[str]) -> str:
     """Generates a safe-to-use filename for uploads.
 
     If possible, tries to keep parts of the original filename intact, such as
     the extension.
 
     :param filename: The original filename.
-    :type filename: str
     :return: The generated filename.
-    :rtype: str
     """
     if filename:
         good_name = secure_filename(filename)
@@ -42,22 +42,19 @@ class DataManager:
     used to access track's images and other on-disk data.
 
     :ivar data_dir: Path to the data folder.
-    :vartype data_dir: pathlib.Path
     """
 
-    def __init__(self, data_dir):
-        self.data_dir = data_dir
+    def __init__(self, data_dir: Path):
+        self.data_dir: Path = data_dir
 
     def _track_data_dir(self, track_id):
         return self.data_dir / "tracks" / str(track_id)
 
-    def images(self, track_id):
+    def images(self, track_id: int) -> List[str]:
         """Returns a list of images that belong to a track.
 
         :param track_id: Numerical ID of the track.
-        :type track_id: int
         :return: A list of image IDs.
-        :rtype: list[str]
         """
         image_dir = self._track_data_dir(track_id) / "images"
         if not image_dir.exists():
@@ -67,14 +64,13 @@ class DataManager:
             images.append(image.name)
         return images
 
-    def purge(self, track_id):
+    def purge(self, track_id: int):
         """Purge all data pertaining to the given track.
 
         This function logs errors but raises no exception, as such it can
         always be used to clean up after a track.
 
         :param track_id: The ID of the track.
-        :type track_id: int
         """
 
         def log_error(_, path, exc_info):
@@ -84,33 +80,26 @@ class DataManager:
         if path.is_dir():
             shutil.rmtree(path, ignore_errors=False, onerror=log_error)
 
-    def image_path(self, track_id, image_id):
+    def image_path(self, track_id: int, image_id: str) -> Path:
         """Returns a path to a saved image.
 
         :raises FileNotFoundError: If the given image could not be found.
         :param track_id: ID of the track.
-        :type track_id: int
         :param image_id: ID of the image.
-        :type image_id: str
         :return: A path pointing to the requested image.
-        :rtype: pathlib.Path
         """
         image = self._track_data_dir(track_id) / "images" / secure_filename(image_id)
         if not image.exists():
             raise FileNotFoundError("The requested image does not exist")
         return image
 
-    def add_image(self, track_id, image, filename=None):
+    def add_image(self, track_id: int, image: BinaryIO, filename: Optional[str] = None) -> str:
         """Saves an image to a track.
 
         :param track_id: ID of the track.
-        :type track_id: int
         :param image: The image, as a file-like object to read from.
-        :type image: file
         :param filename: The image's original filename.
-        :type filename: str
         :return: The ID of the saved image.
-        :rtype: str
         """
         image_dir = self._track_data_dir(track_id) / "images"
         image_dir.mkdir(parents=True, exist_ok=True)
@@ -122,14 +111,12 @@ class DataManager:
 
         return filename
 
-    def delete_image(self, track_id, image_id):
+    def delete_image(self, track_id: int, image_id: str):
         """Deletes an image from a track.
 
         :raises FileNotFoundError: If the given image could not be found.
         :param track_id: ID of the track.
-        :type track_id: int
         :param image_id: ID of the image.
-        :type image_id: str
         """
         # Be sure to not delete anything else than the image file
         image_id = secure_filename(image_id)