make it work with symlinks and imporved format

1 files changed, 17 insertions, 6 deletions

diff --git a/main.py b/main.py
index b9cce95..849f8f9 100644..100755
--- a/main.py
+++ b/main.py
@@ -1,19 +1,30 @@
 #!/usr/bin/env python
 
 from sys import argv, executable
-from os import execv
+from os import execv, path
+
+if path.islink(__file__):
+    from os import readlink
+    file = readlink(__file__)
+else:
+    file = __file__
+cwd = path.dirname(path.abspath(file))
+
+def write():
+    execv(executable, ['python3', f'{cwd}/write.py'])
+def read():
+    execv(executable, ['python3', f'{cwd}/read.py'])
 if len(argv) > 1:
     if argv[1] == 'write':
-        print(' '.join(argv[2:]))
-        execv(executable, ['python3', 'write.py'] + argv[2:])
+        write()
     elif argv[1] =='read':
-        execv(executable, ['python3','read.py'])
+        read()
     else:
         raise Exception("Unknown command\n\tAvailable commands: write, read")
 else:
     from simple_term_menu import TerminalMenu
     choice = TerminalMenu(['Write', 'Read']).show()
     if choice == 0:
-        execv(executable, ['python3', 'write.py'])
+        write()
     elif choice == 1:
-        execv(executable, ['python3','read.py'])
\ No newline at end of file
+        read()