diff --git a/trains_agent/helper/package/translator.py b/trains_agent/helper/package/translator.py index 4d594e2..12f3c4e 100644 --- a/trains_agent/helper/package/translator.py +++ b/trains_agent/helper/package/translator.py @@ -23,6 +23,7 @@ class RequirementsTranslator(object): Path(self.cache_dir).mkdir(parents=True, exist_ok=True) self.config = Config() self.pip = SystemPip(interpreter=interpreter, session=self._session) + self._translate_back = {} def download(self, url): self.pip.download_package(url, cache_dir=self.cache_dir) @@ -60,4 +61,23 @@ class RequirementsTranslator(object): except Exception: command.error('Could not download wheel name of "{}"'.format(line)) return line + + self._translate_back[str(downloaded)] = line + return downloaded + + def replace_back(self, requirements): + if 'pip' in requirements: + original_requirements = requirements['pip'] + new_requirements = [] + for line in original_requirements: + local_file = [d for d in self._translate_back.keys() if d in line] + if local_file: + local_file = local_file[0] + new_requirements.append(line.replace(local_file, self._translate_back[local_file])) + else: + new_requirements.append(line) + + requirements['pip'] = new_requirements + + return requirements