[GitHub] [airflow] mik-laj commented on a change in pull request #6387: [AIRFLOW-5717] Add get_tree_map method to SFTPHook

[GitBox <gi...@apache.org>]

mik-laj commented on a change in pull request #6387: [AIRFLOW-5717] Add get_tree_map method to SFTPHook
URL: https://github.com/apache/airflow/pull/6387#discussion_r337967019
 
 

 ##########
 File path: airflow/contrib/hooks/sftp_hook.py
 ##########
 @@ -232,3 +232,58 @@ def path_exists(self, path):
         """
         conn = self.get_conn()
         return conn.exists(path)
+
+    @staticmethod
+    def _is_path_match(path: str, prefix: Optional[str] = None, delimiter: Optional[str] = None) -> bool:
+        """
+        Return True if given path starts with prefix (if set) and ends with delimiter (if set).
+
+        :param path: path to be checked
+        :type path: str
+        :param prefix: if set path will be checked is starting with prefix
+        :type prefix: str
+        :param delimiter: if set path will be checked is ending with suffix
+        :type delimiter: str
+        :return: bool
+        """
+        if prefix is not None and not path.startswith(prefix):
+            return False
+        if delimiter is not None and not path.endswith(delimiter):
+            return False
+        return True
+
+    def get_tree_map(
+        self, path: str, prefix: Optional[str] = None, delimiter: Optional[str] = None
+    ) -> Dict[str, List[str]]:
+        """
+        Return dictionary with recursive lists of files, directories and unknown paths from given path.
+        It is possible to filter results by giving prefix and/or delimiter parameters.
+
+        :param path: path from which tree will be built
+        :type path: str
+        :param prefix: if set paths will be added if start with prefix
+        :type prefix: str
+        :param delimiter: if set paths will be added if end with delimiter
+        :type delimiter: str
+        :return: dictionary ``{"files": List[str], "dirs": List[str], "unknowns": List[str]}``
+        :rtype: Dict[str, List[str]]
+        """
+        conn = self.get_conn()
+        files, dirs, unknowns = [], [], []
+
+        conn.walktree(
+            remotepath=path,
+            fcallback=lambda x: files.append(x)  # type: ignore
+            if self._is_path_match(x, prefix, delimiter)
+            else None,
+            dcallback=lambda x: dirs.append(x)  # type: ignore
+            if self._is_path_match(x, prefix, delimiter)
+            else None,
+            ucallback=lambda x: unknowns.append(x)  # type: ignore
+            if self._is_path_match(x, prefix, delimiter)
+            else None,
+            recurse=True,
+        )
+
+        result = {"files": files, "dirs": dirs, "unknowns": unknowns}
+        return result
 
 Review comment:
   ```suggestion
           return files, dirs, unknowns
   ```
   This is a more Python approach.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services