Merge pull request #678 from archlinux/torxed-fix-175

Adding support for swap on zram

6 files changed, 27 insertions, 13 deletions

diff --git a/archinstall/lib/disk/blockdevice.py b/archinstall/lib/disk/blockdevice.py
index fad04c7b..f80f57f3 100644
--- a/archinstall/lib/disk/blockdevice.py
+++ b/archinstall/lib/disk/blockdevice.py
@@ -1,10 +1,7 @@
-# flake8: noqa
-# The above ignore, see issue: https://github.com/archlinux/archinstall/pull/650#issuecomment-961995949
 import os
 import json
 import logging
-from .exceptions import DiskError
-from .helpers import all_disks
+from ..exceptions import DiskError
 from ..output import log
 from ..general import SysCommand
 
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index 69593d08..ac3c9d17 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -1,9 +1,9 @@
 import time
 import logging
 import json
-from .exceptions import DiskError
 from .partition import Partition
 from .validators import valid_fs_type
+from ..exceptions import DiskError
 from ..general import SysCommand
 from ..output import log
 from ..storage import storage
diff --git a/archinstall/lib/disk/partition.py b/archinstall/lib/disk/partition.py
index afe46db3..a459a820 100644
--- a/archinstall/lib/disk/partition.py
+++ b/archinstall/lib/disk/partition.py
@@ -7,9 +7,9 @@ import os
 import hashlib
 from typing import Optional
 from .blockdevice import BlockDevice
-from .exceptions import DiskError, SysCallError, UnknownFilesystemFormat
 from .helpers import get_mount_info, get_filesystem_type
-from .storage import storage
+from ..storage import storage
+from ..exceptions import DiskError, SysCallError, UnknownFilesystemFormat
 from ..output import log
 from ..general import SysCommand
 
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 11a0f24e..b400e741 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -6,7 +6,7 @@ import shlex
 import pathlib
 import subprocess
 import glob
-from .disk import get_partitions_in_use, Partition, find_partition
+from .disk import get_partitions_in_use, Partition, find_partition_by_mountpoint
 from .general import SysCommand
 from .hardware import has_uefi, is_vm, cpu_vendor
 from .locale_helpers import verify_keyboard_layout, verify_x11_keyboard_layout
@@ -14,7 +14,6 @@ from .disk.helpers import get_mount_info
 from .mirrors import use_mirrors
 from .plugins import plugins
 from .storage import storage
-from .systemd import Boot
 # from .user_interaction import *
 from .output import log
 from .profiles import Profile
@@ -270,6 +269,7 @@ class Installer:
 			fh.write("NTP=0.arch.pool.ntp.org 1.arch.pool.ntp.org 2.arch.pool.ntp.org 3.arch.pool.ntp.org\n")
 			fh.write("FallbackNTP=0.pool.ntp.org 1.pool.ntp.org 0.fr.pool.ntp.org\n")
 
+		from .systemd import Boot
 		with Boot(self) as session:
 			session.SysCommand(["timedatectl", "set-ntp", 'true'])
 
@@ -463,6 +463,22 @@ class Installer:
 
 		return True
 
+	def setup_swap(self, kind='zram'):
+		if kind == 'zram':
+			self.log(f"Setting up swap on zram")
+			self.pacstrap('zram-generator')
+			
+			# We could use the default example below, but maybe not the best idea: https://github.com/archlinux/archinstall/pull/678#issuecomment-962124813
+			# zram_example_location = '/usr/share/doc/zram-generator/zram-generator.conf.example'
+			# shutil.copy2(f"{self.target}{zram_example_location}", f"{self.target}/usr/lib/systemd/zram-generator.conf")
+			with open(f"{self.target}/etc/systemd/zram-generator.conf", "w") as zram_conf:
+				zram_conf.write("[zram0]\n")
+
+			if self.enable_service('systemd-zram-setup@zram0.service'):
+				return True
+		else:
+			raise ValueError(f"Archinstall currently only supports setting up swap on zram")
+
 	def add_bootloader(self, bootloader='systemd-bootctl'):
 		for plugin in plugins.values():
 			if hasattr(plugin, 'on_add_bootloader'):
@@ -575,7 +591,7 @@ class Installer:
 				self.helper_flags['bootloader'] = True
 				return True
 			else:
-				boot_partition = find_partition(mountpoint=f"{self.target}/boot")
+				boot_partition = find_partition_by_mountpoint(self.partitions, relative_mountpoint=f"/boot")
 				SysCommand(f'/usr/bin/arch-chroot {self.target} grub-install --target=i386-pc --recheck {boot_partition.path}')
 				SysCommand(f'/usr/bin/arch-chroot {self.target} grub-mkconfig -o /boot/grub/grub.cfg')
 				self.helper_flags['bootloader'] = True
@@ -651,7 +667,6 @@ class Installer:
 			# In accordance with https://github.com/archlinux/archinstall/issues/107#issuecomment-841701968
 			# Setting an empty keymap first, allows the subsequent call to set layout for both console and x11.
 			from .systemd import Boot
-
 			with Boot(self) as session:
 				session.SysCommand(["localectl", "set-keymap", '""'])
 
@@ -675,7 +690,6 @@ class Installer:
 				return False
 
 			from .systemd import Boot
-
 			with Boot(self) as session:
 				session.SysCommand(["localectl", "set-x11-keymap", '""'])
 
diff --git a/archinstall/lib/systemd.py b/archinstall/lib/systemd.py
index d297c507..a7e35839 100644
--- a/archinstall/lib/systemd.py
+++ b/archinstall/lib/systemd.py
@@ -1,5 +1,4 @@
 import logging
-
 from .general import SysCommand, SysCommandWorker, locate_binary
 from .installer import Installer
 from .output import log
diff --git a/archinstall/lib/user_interaction.py b/archinstall/lib/user_interaction.py
index 6b142288..e62d8c6c 100644
--- a/archinstall/lib/user_interaction.py
+++ b/archinstall/lib/user_interaction.py
@@ -325,6 +325,10 @@ class MiniCurses:
 			return response
 
 
+def ask_for_swap(prompt='Would you like to use swap on zram? (Y/n): ', forced=False):
+	return True if input(prompt).strip(' ').lower() not in ('n', 'no') else False
+
+
 def ask_for_superuser_account(prompt='Username for required superuser with sudo privileges: ', forced=False):
 	while 1:
 		new_user = input(prompt).strip(' ')