From c2c57b47d52de1216b1af314606bf48d2eaa9a6e Mon Sep 17 00:00:00 2001 From: kris Date: Fri, 14 Jun 2019 22:03:37 +0100 Subject: [PATCH] Add type annotations --- transcoder/screen.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/transcoder/screen.py b/transcoder/screen.py index 956860e..1fe68a7 100644 --- a/transcoder/screen.py +++ b/transcoder/screen.py @@ -2,9 +2,13 @@ import functools import pickle +from typing import Union import numpy as np +# Type annotation for cases where we may process either an int or a numpy array. +IntOrArray = Union[int, np.ndarray] + def y_to_base_addr(y: int, page: int = 0) -> int: """Maps y coordinate to base address on given screen page""" @@ -147,7 +151,7 @@ class DHGRBitmap: self.packed = np.empty(shape=(32, 128), dtype=np.uint32) self._pack() - def _pack(self): + def _pack(self) -> None: """Interleave and pack aux and main memory into 28-bit uint32 array""" # Palette bit is unused for DHGR so mask it out @@ -180,7 +184,10 @@ class DHGRBitmap: return 1 @staticmethod - def masked_update(byte_offset: int, old_value, new_value: int): + def masked_update( + byte_offset: int, + old_value: IntOrArray, + new_value: int) -> IntOrArray: # Mask out 7-bit value where update will go masked_value = old_value & ~(0x7f << (7 * byte_offset)) @@ -188,7 +195,7 @@ class DHGRBitmap: return masked_value ^ update - def apply(self, page: int, offset: int, is_aux: bool, value: int): + def apply(self, page: int, offset: int, is_aux: bool, value: int) -> None: """Update packed representation of changing main/aux memory.""" byte_offset = self.interleaved_byte_offset(offset, is_aux) @@ -197,7 +204,10 @@ class DHGRBitmap: self.packed[page, packed_offset] = self.masked_update( byte_offset, self.packed[page, packed_offset], value) - def mask_and_shift_data(self, data, byte_offset): + def mask_and_shift_data( + self, + data: IntOrArray, + byte_offset: int) -> IntOrArray: """Masks and shifts data into the 8 or 12-bit range.""" return (data & self.BYTE_MASK32[byte_offset]) >> ( self.BYTE_SHIFTS[byte_offset])