From b4eb8557f5ab4ced63cae061875cf6080f0ab7cd Mon Sep 17 00:00:00 2001
From: Anton Hvornum <anton.feeds@gmail.com>
Date: Mon, 1 Nov 2021 09:59:32 +0000
Subject: Adding support for swap on zram

---
 archinstall/lib/installer.py        | 11 +++++++++++
 archinstall/lib/user_interaction.py |  3 +++
 examples/guided.py                  |  4 ++++
 3 files changed, 18 insertions(+)

diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 0bdddb2e..f83e0128 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -1,4 +1,5 @@
 import time
+import shutil
 from .disk import *
 from .hardware import *
 from .locale_helpers import verify_keyboard_layout, verify_x11_keyboard_layout
@@ -9,6 +10,7 @@ from .storage import storage
 from .user_interaction import *
 from .disk.btrfs import create_subvolume, mount_subvolume
 from .exceptions import DiskError, ServiceException
+from .output import log
 
 # Any package that the Installer() is responsible for (optional and the default ones)
 __packages__ = ["base", "base-devel", "linux-firmware", "linux", "linux-lts", "linux-zen", "linux-hardened"]
@@ -442,6 +444,15 @@ 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')
+			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")
+		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'):
diff --git a/archinstall/lib/user_interaction.py b/archinstall/lib/user_interaction.py
index ba6259b1..d0ed0747 100644
--- a/archinstall/lib/user_interaction.py
+++ b/archinstall/lib/user_interaction.py
@@ -326,6 +326,9 @@ 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(' ')
diff --git a/examples/guided.py b/examples/guided.py
index 60539a0b..47d1589b 100644
--- a/examples/guided.py
+++ b/examples/guided.py
@@ -132,6 +132,8 @@ def ask_user_questions():
 	if not archinstall.arguments.get("bootloader", None):
 		archinstall.arguments["bootloader"] = archinstall.ask_for_bootloader()
 
+	if not archinstall.arguments.get('swap', None):
+		archinstall.archinstall['swap'] = archinstall.ask_for_swap()
 
 	# Get the hostname for the machine
 	if not archinstall.arguments.get('hostname', None):
@@ -282,6 +284,8 @@ def perform_installation(mountpoint):
 			if archinstall.arguments["bootloader"] == "grub-install" and has_uefi():
 				installation.add_additional_packages("grub")
 			installation.add_bootloader(archinstall.arguments["bootloader"])
+			if archinstall.archinstall['swap']:
+				installation.setup_swap('zram')
 
 			# If user selected to copy the current ISO network configuration
 			# Perform a copy of the config
-- 
cgit v1.2.3-70-g09d2