Speed up list_available (#133)

* Use mamba repoquery to retrieve packages instead of conda search when mamba is available to speed up the list_available method.

* Avoid duplicated code

* Simplify processing of mamba repoquery output

* Add testing for mamba repoquery

Co-authored-by: Frederic COLLONVAL <[email protected]>

Author: ericpre

mamba_gator/envmanager.py

@@ -1,14 +1,14 @@
 # Copyright (c) 2016-2020 Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 import asyncio
+import collections
 import json
 import logging
 import os
 import re
-import ssl
 import sys
 import tempfile
-from functools import partial
+from functools import partial, lru_cache
 from pathlib import Path
 from subprocess import PIPE, Popen
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -233,6 +233,10 @@ def manager(self) -> str:
 
         return EnvManager._manager_exe
 
+    @lru_cache()
+    def is_mamba(self):
+        return Path(self.manager).stem == "mamba"
+
     async def env_channels(
         self, configuration: Optional[Dict[str, Any]] = None
     ) -> Dict[str, Dict[str, List[str]]]:
@@ -552,7 +556,7 @@ async def pkg_depends(self, pkg: str) -> Dict[str, List[str]]:
         Returns:
             {"package": List[dependencies]}
         """
-        if Path(self.manager).stem != "mamba":
+        if not self.is_mamba():
             self.log.warning(
                 "Package manager '{}' does not support dependency query.".format(
                     self.manager
@@ -585,7 +589,10 @@ async def list_available(self) -> Dict[str, List[Dict[str, str]]]:
                 "with_description": bool  # Whether we succeed in get some channeldata.json files
             }
         """
-        ans = await self._execute(self.manager, "search", "--json")
+        if self.is_mamba():
+            ans = await self._execute(self.manager, "repoquery", "search", "*", "--json")
+        else:
+            ans = await self._execute(self.manager, "search", "--json")
         _, output = ans
 
         current_loop = tornado.ioloop.IOLoop.current()
@@ -596,6 +603,23 @@ async def list_available(self) -> Dict[str, List[Dict[str, str]]]:
             # dictionary with error info
             return data
 
+        def process_mamba_repoquery_output(data: Dict) -> Dict:
+            """Make a dictionary with keys as packages name and values
+            containing the list of available packages to match the json output
+            of "conda search --json".
+            """
+
+            data_ = collections.defaultdict(lambda : [])
+            for entry in data['result']['pkgs']:
+                name = entry.get('name')
+                if name is not None:
+                    data_[name].append(entry)
+
+            return data_
+
+        if self.is_mamba():
+            data = await current_loop.run_in_executor(None, process_mamba_repoquery_output, data)
+
         def format_packages(data: Dict) -> List:
             packages = []
 

mamba_gator/tests/test_api.py

@@ -5,6 +5,7 @@
 import sys
 import unittest
 import unittest.mock as mock
+from itertools import chain
 
 try:
     from unittest.mock import AsyncMock
@@ -971,6 +972,14 @@ def test_package_list_available(self):
                     "ssl_verify": False,
                 }
 
+                if has_mamba:
+                    # Change dummy to match mamba repoquery format
+                    dummy = {
+                        "result": {
+                            "pkgs": list(chain(*dummy.values()))
+                        }
+                    }
+
                 rvalue = [
                     (0, json.dumps(dummy)),
                     (0, json.dumps(channels)),
@@ -984,6 +993,13 @@ def test_package_list_available(self):
 
                 r = self.wait_for_task(self.conda_api.get, ["packages"])
                 self.assertEqual(r.status_code, 200)
+
+                args, _ = f.call_args_list[0]
+                if has_mamba:
+                    self.assertSequenceEqual(args[1:], ["repoquery", "search", "*", "--json"])
+                else:
+                    self.assertSequenceEqual(args[1:], ["search", "--json"])
+
                 body = r.json()
 
                 expected = {
@@ -1164,6 +1180,14 @@ def test_package_list_available_local_channel(self):
                     ],
                 }
 
+                if has_mamba:
+                    # Change dummy to match mamba repoquery format
+                    dummy = {
+                        "result": {
+                            "pkgs": list(chain(*dummy.values()))
+                        }
+                    }
+
                 with tempfile.TemporaryDirectory() as local_channel:
                     with open(
                         os.path.join(local_channel, "channeldata.json"), "w+"
@@ -1202,6 +1226,13 @@ def test_package_list_available_local_channel(self):
 
                     r = self.wait_for_task(self.conda_api.get, ["packages"])
                     self.assertEqual(r.status_code, 200)
+
+                    args, _ = f.call_args_list[0]
+                    if has_mamba:
+                        self.assertSequenceEqual(args[1:], ["repoquery", "search", "*", "--json"])
+                    else:
+                        self.assertSequenceEqual(args[1:], ["search", "--json"])
+
                     body = r.json()
 
                     expected = {
@@ -1382,6 +1413,14 @@ def test_package_list_available_no_description(self):
                     ],
                 }
 
+                if has_mamba:
+                    # Change dummy to match mamba repoquery format
+                    dummy = {
+                        "result": {
+                            "pkgs": list(chain(*dummy.values()))
+                        }
+                    }
+
                 with tempfile.TemporaryDirectory() as local_channel:
                     local_name = local_channel.strip("/")
                     channels = {
@@ -1414,6 +1453,13 @@ def test_package_list_available_no_description(self):
 
                     r = self.wait_for_task(self.conda_api.get, ["packages"])
                     self.assertEqual(r.status_code, 200)
+
+                    args, _ = f.call_args_list[0]
+                    if has_mamba:
+                        self.assertSequenceEqual(args[1:], ["repoquery", "search", "*", "--json"])
+                    else:
+                        self.assertSequenceEqual(args[1:], ["search", "--json"])
+
                     body = r.json()
 
                     expected = {
@@ -1620,6 +1666,15 @@ def test_package_list_available_caching(self):
                     "ssl_verify": False,
                 }
 
+
+                if has_mamba:
+                    # Change dummy to match mamba repoquery format
+                    dummy = {
+                        "result": {
+                            "pkgs": list(chain(*dummy.values()))
+                        }
+                    }
+
                 rvalue = [
                     (0, json.dumps(dummy)),
                     (0, json.dumps(channels)),
@@ -1635,6 +1690,12 @@ def test_package_list_available_caching(self):
                 r = self.wait_for_task(self.conda_api.get, ["packages"])
                 self.assertEqual(r.status_code, 200)
 
+                args, _ = f.call_args_list[0]
+                if has_mamba:
+                    self.assertSequenceEqual(args[1:], ["repoquery", "search", "*", "--json"])
+                else:
+                    self.assertSequenceEqual(args[1:], ["search", "--json"])
+
                 expected = {
                     "packages": [
                         {