diff --git a/colabcode/code.py b/colabcode/code.py index 56015ca..b10c9c3 100644 --- a/colabcode/code.py +++ b/colabcode/code.py @@ -15,14 +15,15 @@ class ColabCode: - def __init__(self, port=10000, password=None, authtoken=None, mount_drive=False): + def __init__(self, port=10000, password=None, authtoken=None, mount_drive=False, drive_path=None): self.port = port self.password = password self.authtoken = authtoken self._mount = mount_drive + self._default_path = "/content/drive" self._install_code() self._install_extensions() - self._start_server() + self._start_server(drive_path) self._run_code() def _install_code(self): @@ -34,11 +35,14 @@ def _install_code(self): stdout=subprocess.PIPE, ) + def _install_extensions(self): for ext in EXTENSIONS: subprocess.run(["code-server", "--install-extension", f"{ext}"]) - def _start_server(self): + + + def _start_server(self, drive_path): if self.authtoken: ngrok.set_auth_token(self.authtoken) active_tunnels = ngrok.get_tunnels() @@ -46,12 +50,26 @@ def _start_server(self): public_url = tunnel.public_url ngrok.disconnect(public_url) url = ngrok.connect(addr=self.port, options={"bind_tls": True}) - print(f"Code Server can be accessed on: {url}") + print(self._mount_drive(url, drive_path)) + + + def _mount_drive(self, url, drive_path): + if self._mount and colab_env: + if drive_path: + try : + drive.mount(drive_path) + self._default_path = drive_path + except: + drive.mount(self._default_path) + else: + drive.mount(self._default_path) + return f"Code server can be accessed on : {url}/?folder={self._default_path}" + return f"Code server can be accessed on : {url}" + + def _run_code(self): os.system(f"fuser -n tcp -k {self.port}") - if self._mount and colab_env: - drive.mount("/content/drive") if self.password: code_cmd = f"PASSWORD={self.password} code-server --port {self.port} --disable-telemetry" else: diff --git a/scripts/colabcode b/scripts/colabcode index 6b7354b..a2bb6fe 100644 --- a/scripts/colabcode +++ b/scripts/colabcode @@ -28,6 +28,12 @@ if __name__ == "__main__": action="store_true", help="if you use --mount_drive, your google drive will be mounted", ) + optional.add_argument( + "--drive_absolute_path", + type=str, + help="open your code-server directly from the folder you want. Accepts absolute path.", + default=None + ) args = parser.parse_args()