diff --git a/README.md b/README.md index 2c80312..ddca1cb 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,8 @@ c.DrivesConfig.access_key_id = "" c.DrivesConfig.secret_access_key = "" c.DrivesConfig.session_token = "" c.DrivesConfig.provider = "" +c.DrivesConfig.included_drives = "" +c.DrivesConfig.excluded_drives = "" ``` ### Custom credentials file path @@ -76,6 +78,8 @@ export JP_DRIVES_ACCESS_KEY_ID="" export JP_DRIVES_SECRET_ACCESS_KEY="" export JP_DRIVES_SESSION_TOKEN="" export JP_DRIVES_CUSTOM_CREDENTIALS_PATH="" +export JP_DRIVES_INCLUDED_DRIVES ="" +export JP_DRIVES_EXCLUDED_DRIVES="" ``` ## Uninstall diff --git a/jupyter_drives/base.py b/jupyter_drives/base.py index 32f536b..43343ba 100644 --- a/jupyter_drives/base.py +++ b/jupyter_drives/base.py @@ -1,7 +1,7 @@ import os from sys import platform import entrypoints -from traitlets import Enum, Unicode, default +from traitlets import Enum, Unicode, default, Set from traitlets.config import Configurable import boto3 @@ -71,10 +71,27 @@ def set_default_api_base_url(self): help="The source control provider.", ) + excluded_drives = Set( + trait=Unicode(), + config=True, + help="List of drives that should be excluded from drive browser listing. Drive names should be separated by spaces.", + ) + + included_drives = Set( + trait=Unicode(), + config=True, + help="List of drives that should be included in drive browser listing. Drive names should be separated by spaces.", + ) + def __init__(self, **kwargs): super().__init__(**kwargs) # check if credentials were already set in jupyter_notebook_config.py self.credentials_already_set = self.access_key_id is not None and self.secret_access_key is not None + + # check list of excluded and included drives + self.check_excluded_and_included_drives() + + # load credentials self.load_credentials() def load_credentials(self): @@ -111,4 +128,26 @@ def load_credentials(self): self.access_key_id = c.access_key self.secret_access_key = c.secret_key self.region_name = s.region_name - self.session_token = c.token \ No newline at end of file + self.session_token = c.token + + def check_excluded_and_included_drives(self): + # list of drives to exclude was provided + if len(self.excluded_drives) != 0: + return + + # list of what drives to include was provided + if len(self.included_drives) != 0: + return + + # check environment variables for excluded or included drives configuration, only one is taken into account + if "JP_DRIVES_EXCLUDED_DRIVES" in os.environ: + drives = os.environ["JP_DRIVES_EXCLUDED_DRIVES"] + self.excluded_drives = drives.split(' ') + return + + if "JP_DRIVES_INCLUDED_DRIVES" in os.environ: + drives = os.environ["JP_DRIVES_INCLUDED_DRIVES"] + self.included_drives = drives.split(' ') + return + + return diff --git a/jupyter_drives/manager.py b/jupyter_drives/manager.py index 1348338..f4fc6ca 100644 --- a/jupyter_drives/manager.py +++ b/jupyter_drives/manager.py @@ -53,8 +53,9 @@ def __init__(self, config: traitlets.config.Config) -> None: self._max_files_listed = 1025 self._drives = None self._external_drives = {} - self._excluded_drives = set() - + self._excluded_drives = self._config.excluded_drives if len(self._config.excluded_drives) != 0 else set() + self._included_drives = self._config.included_drives if len(self._config.included_drives) != 0 else set() + # instate fsspec file system self._file_system = fsspec.filesystem(self._config.provider, asynchronous=True) @@ -255,7 +256,14 @@ async def list_drives(self): reason=f"The following error occured when listing drives: {e}", ) - for result in results: + if len(self._included_drives) != 0: + for result in results: + if result.name not in self._included_drives: + self._excluded_drives.add(result.name) + # clear list once initialized + self._included_drives.clear() + + for result in results: if result.name not in self._excluded_drives: data.append( {