Add type hintings to overlay methods

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>

subiquity/server/apt.py

@@ -18,6 +18,7 @@ import logging
 import os
 import shutil
 import tempfile
+from typing import List, Optional, Union
 
 from curtin.config import merge_config
 
@@ -32,7 +33,7 @@ log = logging.getLogger('subiquity.server.apt')
 
 class _MountBase:
 
-    def p(self, *args):
+    def p(self, *args: str) -> str:
         for a in args:
             if a.startswith('/'):
                 raise Exception('no absolute paths here please')
@@ -44,15 +45,20 @@ class _MountBase:
 
 
 class Mountpoint(_MountBase):
-    def __init__(self, *, mountpoint):
-        self.mountpoint = mountpoint
+    def __init__(self, *, mountpoint: str):
+        self.mountpoint: str = mountpoint
 
 
 class OverlayMountpoint(_MountBase):
-    def __init__(self, *, lowers, upperdir, mountpoint):
-        self.lowers = lowers
-        self.upperdir = upperdir
-        self.mountpoint = mountpoint
+    def __init__(self, *, lowers, upperdir: Optional[str], mountpoint: str):
+        # The first element in lowers will be the bottom layer and the last
+        # element will be the top layer.
+        self.lowers: List[Lower] = lowers
+        self.upperdir: Optional[str] = upperdir
+        self.mountpoint: str = mountpoint
+
+
+Lower = Union[Mountpoint, str, OverlayMountpoint]
 
 
 @functools.singledispatch
@@ -119,10 +125,11 @@ class AptConfigurer:
     #    system, or if it is not, just copy /var/lib/apt/lists from the
     #    'configured_tree' overlay.
 
-    def __init__(self, app, source):
+    def __init__(self, app, source: str):
         self.app = app
-        self.source = source
-        self.configured_tree = None
+        self.source: str = source
+        self.configured_tree: Optional[OverlayMountpoint] = None
+        self.install_tree: Optional[OverlayMountpoint] = None
         self.install_mount = None
         self._mounts = []
         self._tdirs = []
@@ -144,10 +151,11 @@ class AptConfigurer:
         self._mounts.append(m)
         return m
 
-    async def unmount(self, mountpoint):
+    async def unmount(self, mountpoint: str):
         await self.app.command_runner.run(['umount', mountpoint])
 
-    async def setup_overlay(self, lowers):
+    async def setup_overlay(self, lowers: Union[List[Lower], Lower]) \
+            -> OverlayMountpoint:
         tdir = self.tdir()
         target = f'{tdir}/mount'
         lowerdir = lowerdir_for(lowers)
@@ -221,23 +229,23 @@ class AptConfigurer:
         for d in self._tdirs:
             shutil.rmtree(d)
 
-    async def deconfigure(self, context, target):
-        target = Mountpoint(mountpoint=target)
+    async def deconfigure(self, context, target: str) -> None:
+        target_mnt = Mountpoint(mountpoint=target)
 
         async def _restore_dir(dir):
-            shutil.rmtree(target.p(dir))
+            shutil.rmtree(target_mnt.p(dir))
             await self.app.command_runner.run([
-                'cp', '-aT', self.configured_tree.p(dir), target.p(dir),
+                'cp', '-aT', self.configured_tree.p(dir), target_mnt.p(dir),
                 ])
 
-        await self.unmount(target.p('cdrom'))
-        os.rmdir(target.p('cdrom'))
+        await self.unmount(target_mnt.p('cdrom'))
+        os.rmdir(target_mnt.p('cdrom'))
 
         await _restore_dir('etc/apt')
 
         if self.app.base_model.network.has_network:
             await run_curtin_command(
-                self.app, context, "in-target", "-t", target.p(),
+                self.app, context, "in-target", "-t", target_mnt.p(),
                 "--", "apt-get", "update")
         else:
             await _restore_dir('var/lib/apt/lists')
@@ -246,15 +254,22 @@ class AptConfigurer:
 
         if self.app.base_model.network.has_network:
             await run_curtin_command(
-                self.app, context, "in-target", "-t", target.p(),
+                self.app, context, "in-target", "-t", target_mnt.p(),
                 "--", "apt-get", "update")
 
 
 class DryRunAptConfigurer(AptConfigurer):
 
-    async def setup_overlay(self, source):
-        if isinstance(source, OverlayMountpoint):
-            source = source.lowers[0]
+    async def setup_overlay(self, lowers: Union[List[Lower], Lower]) \
+            -> OverlayMountpoint:
+        # XXX This implementation expects that:
+        # - on first invocation, we pass a string as "lowers"
+        # - on second invocation, we pass the OverlayMountPoint returned by the
+        # first invocation.
+        if isinstance(lowers, OverlayMountpoint):
+            source = lowers.lowers[0]
+        else:
+            source = lowers
         target = self.tdir()
         os.mkdir(f'{target}/etc')
         await arun_command([
@@ -272,7 +287,7 @@ class DryRunAptConfigurer(AptConfigurer):
         return
 
 
-def get_apt_configurer(app, source):
+def get_apt_configurer(app, source: str):
     if app.opts.dry_run:
         return DryRunAptConfigurer(app, source)
     else:

subiquity/server/controllers/source.py

@@ -13,6 +13,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+from typing import Optional
 import os
 
 from curtin.commands.extract import get_handler_for_source
@@ -54,7 +55,7 @@ class SourceController(SubiquityController):
     def __init__(self, app):
         super().__init__(app)
         self._handler = None
-        self.source_path = None
+        self.source_path: Optional[str] = None
 
     def start(self):
         path = '/cdrom/casper/install-sources.yaml'