diff --git a/modules/swap/__init__.py b/modules/swap/__init__.py
new file mode 100644
index 0000000..9e23e02
--- /dev/null
+++ b/modules/swap/__init__.py
@@ -0,0 +1,5 @@
+from .module import SwapModule
+
+
+def get_module():
+    return SwapModule()
\ No newline at end of file
diff --git a/modules/swap/module.py b/modules/swap/module.py
index 519dfe5..0b0dcee 100644
--- a/modules/swap/module.py
+++ b/modules/swap/module.py
@@ -1,61 +1,76 @@
 import subprocess
-from PyInquirer import prompt
 from AbstractModule import AbstractModule
+from PyInquirer import prompt
 
 
-class VimModule(AbstractModule):
+class SwapModule(AbstractModule):
+
     def __init__(self):
         super().__init__()
+        self.swap_file = None
 
+    def run(self):
 
-# Check default swap location
-def convert_string(string):
-    li = list(string.split())
-    return li
+        actions = ["Create/Resize", "Delete"]
+        menu = [
+            {
+                'type': 'list',
+                'message': 'Select action',
+                'name': 'action',
+                'choices': list(map(lambda x: {"name": x}, actions))
+            }
+        ]
+        selected_action = prompt(menu)['action']
+        if selected_action is "Create/Resize":
+            self.create_resize_swap()
+        elif selected_action is "Delete":
+            self.delete_swap()
 
+    # Check default swap location
+    def swap_location_check(self) -> None:
+        output_swaps = subprocess.check_output(['cat', '/proc/swaps']).decode("UTF-8")
+        li = list(output_swaps.split())
+        final_list = li
+        try:
+            print("Swap is located here:", final_list[5])
+            self.swap_file = final_list[5]
+        except IndexError:
+            print("Swap file doesn´t exist!\n")
+    # End of location check
 
-output_swaps = subprocess.check_output(['cat', '/proc/swaps']).decode("UTF-8")
-final_list = convert_string(output_swaps)
-print(final_list[5])
-# End of location check
+    def create_resize_swap(self):
+        output_swapon = subprocess.check_output(['swapon', '--show']).decode("UTF-8")
+        if not output_swapon:
+            size = input("How big should the swap be(numbers in Gigabyte)? ")
+            create_swapfile = """sudo fallocate -l {}G /swapfile && 
+            sudo chmod 600 /swapfile && 
+            sudo mkswap /swapfile && 
+            sudo swapon /swapfile""".format(size)
+            subprocess.call(create_swapfile, shell=True)
 
-output_swapon = subprocess.check_output(['swapon', '--show']).decode("UTF-8")
-output_free = subprocess.check_output(['free', '-h']).decode("UTF-8")
+            output_free = subprocess.check_output(['free', '-h']).decode("UTF-8")
+            print(output_free.strip())
 
-option = input("Enter option. 1.Create/Resize 2.Delete: ")
+        else:
+            print("")
+            print("Swap already installed! You can resize it!")
+            resize = input("How big should your swap become(numbers in Gigabyte)? ")
+            resize_swapfile = """sudo swapoff /swapfile && 
+            sudo fallocate -l {}G /swapfile && 
+            sudo mkswap /swapfile && 
+            sudo swapon /swapfile""".format(resize)
+            subprocess.call(resize_swapfile, shell=True)
+            output_free = subprocess.check_output(['free', '-h']).decode("UTF-8")
+            print(output_free.strip())
 
-if option == "1":
-    if not output_swapon:
-        size = input("How big should the swap be(numbers in Gigabyte)? ")
-        create_swapfile = """sudo fallocate -l {}G /swapfile && 
-        sudo chmod 600 /swapfile && 
-        sudo mkswap /swapfile && 
-        sudo swapon /swapfile""".format(size)
-        subprocess.call(create_swapfile, shell=True)
+    def delete_swap(self):
+        self.swap_location_check()
+        disable_swapfile = """sudo swapoff {0} && 
+        sudo rm {0}""".format(self.swap_file)
 
+        subprocess.call(disable_swapfile, shell=True)
+        output_swapon = subprocess.check_output(['swapon', '--show']).decode("UTF-8")
         output_free = subprocess.check_output(['free', '-h']).decode("UTF-8")
-        print(output_free.strip())
-
-    else:
-        print("")
-        print("Swap already installed! You can resize it!")
-        resize = input("How big should your swap become(numbers in Gigabyte)? ")
-        resize_swapfile = """sudo swapoff /swapfile && 
-        sudo fallocate -l {}G /swapfile && 
-        sudo mkswap /swapfile && 
-        sudo swapon /swapfile""".format(resize)
-        subprocess.call(resize_swapfile, shell=True)
-        output_free = subprocess.check_output(['free', '-h']).decode("UTF-8")
-        print(output_free.strip())
-elif option == "2":
-    disable_swapfile = """sudo swapoff /swapfile && 
-    sudo rm /swapfile"""
-
-    subprocess.call(disable_swapfile, shell=True)
-    output_swapon = subprocess.check_output(['swapon', '--show']).decode("UTF-8")
-    output_free = subprocess.check_output(['free', '-h']).decode("UTF-8")
-    if not output_swapon:
-        print("Swap deleted!")
-        print(output_free.strip())
-else:
-    print("Wrong input!")
+        if not output_swapon:
+            print("Swap deleted!")
+            print(output_free.strip())