simplify os check

organize.py

@@ -9,52 +9,49 @@ FILE_NAME = "organize.py"
 EXT_NAME = "ext.py"
 
 
-def organize_files(drive):
-    # check os
+def organize_files(path):
     operating_system = platform.system()
     if operating_system == "Windows":  # Windows
-        drive = drive + "\\"
+        path = path + "\\"
     elif operating_system == "Darwin":  # Mac
-        drive = os.path.expanduser(drive)
+        path = os.path.expanduser(path)
     else:
         print(f"Operating system {operating_system} not currently supported")
         sys.exit(0)
-    if not os.path.exists(drive):
-        print(f"ERROR! {drive} is not a valid location")
+    if not os.path.exists(path):
+        print("ERROR! Invalid location")
         return
-    files = os.listdir(drive)
-    extns = {os.path.splitext(file)[1].strip(".") for file in files}
+    files = os.listdir(path)
+    extns = {os.path.splitext(file)[1].strip('.') for file in files}
 
     # Create Folders
     for ext in extns:
         folder = foldername(ext)
-        if folder and not os.path.exists(drive + folder):
-            os.makedirs(drive + folder)
+        if folder and not os.path.exists(path + folder):
+            os.makedirs(path + folder)
 
     # Move Files To Folders
     for file in files:
         if file in [FILE_NAME, EXT_NAME]:
             continue
-        ext = os.path.splitext(file)[1].strip(".")
+        ext = os.path.splitext(file)[1].strip('.')
         folder = foldername(ext)
         if not folder:
             continue
 
-        src = drive + file
-        dest = drive + folder + "/" + file
+        src = path + file
+        dest = path + folder + '/' + file
 
         if not os.path.exists(dest):
             shutil.move(src, dest)
             print(f"Moved {file} to {folder}")
 
-        print(f"\nSUCCESS! All files organized in {drive}")
+    print(f"\nSUCCESS! All files organized in {path}")
 
 
-if __name__ == "__main__":
-    try:
-        # location = sys.argv[1]
-        # organize_files(location)
-        organize_files('~/Trash')
-    except Exception as e:
-        print("USAGE: python organize.py <location>")
-        print(e)
+try:
+    location = sys.argv[1]
+    organize_files(location)
+except Exception as e:
+    print(f"error {e}")
+    print("USAGE: python organize.py <location>")