diff --git a/mim/commands/search.py b/mim/commands/search.py index bda992e..790df1a 100644 --- a/mim/commands/search.py +++ b/mim/commands/search.py @@ -4,7 +4,6 @@ import re import subprocess import tempfile import typing -from pkg_resources import resource_filename from typing import Any, List, Optional import click @@ -19,6 +18,7 @@ from mim.utils import ( cast2lowercase, echo_success, get_github_url, + get_installed_path, highlighted_error, is_installed, split_package_version, @@ -212,14 +212,16 @@ def load_metadata_from_local(package: str): if is_installed(package): # rename the model_zoo.yml to model-index.yml but support both of them # for backward compatibility - metadata_path = resource_filename(package, 'model-index.yml') + installed_path = get_installed_path(package) + metadata_path = osp.join(installed_path, 'model-index.yml') if not osp.exists(metadata_path): - metadata_path = resource_filename(package, 'model_zoo.yml') + metadata_path = osp.join(installed_path, 'model_zoo.yml') if not osp.exists(metadata_path): raise FileNotFoundError( highlighted_error( - 'model-index.yml or model_zoo.yml is not found, please' - f' upgrade your {package} to support search command')) + f'{installed_path}/model-index.yml or {installed_path}' + '/model_zoo.yml is not found, please upgrade your ' + f'{package} to support search command')) metadata = load(metadata_path) @@ -265,9 +267,9 @@ def load_metadata_from_remote(package: str) -> Optional[ModelIndex]: # rename the model_zoo.yml to model-index.yml but support both of # them for backward compatibility - metadata_path = resource_filename(package, 'model-index.yml') + metadata_path = osp.join(repo_root, 'model-index.yml') if not osp.exists(metadata_path): - metadata_path = resource_filename(package, 'model_zoo.yml') + metadata_path = osp.join(repo_root, 'model_zoo.yml') if not osp.exists(metadata_path): raise FileNotFoundError( highlighted_error( @@ -277,7 +279,7 @@ def load_metadata_from_remote(package: str) -> Optional[ModelIndex]: metadata = load(metadata_path) - if not version: + if version: with open(pkl_path, 'wb') as fw: pickle.dump(metadata, fw) diff --git a/tests/test_search.py b/tests/test_search.py index 2130b7f..77cf933 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -1,8 +1,10 @@ +import os.path as osp + from click.testing import CliRunner from mim.commands.install import cli as install from mim.commands.search import cli as search -from mim.utils import is_installed +from mim.utils import DEFAULT_CACHE_DIR, is_installed def setup_module(): @@ -23,10 +25,18 @@ def test_search(): # search master branch result = runner.invoke(search, ['mmcls', '--remote']) assert result.exit_code == 0 + # mim search mmsegmentation --remote + result = runner.invoke(search, ['mmsegmentation', '--remote']) + assert result.exit_code == 0 + # mim search mmaction2 --remote + result = runner.invoke(search, ['mmaction2', '--remote']) + assert result.exit_code == 0 # mim search mmcls==0.11.0 --remote result = runner.invoke(search, ['mmcls==0.11.0', '--remote']) assert result.exit_code == 0 + # the metadata of mmcls==0.11.0 will be saved in cache + assert osp.exists(osp.join(DEFAULT_CACHE_DIR, 'mmcls-0.11.0.pkl')) # mim search mmcls --model res # invali model