From 259113c989bd19f8f3aceab35b34655c82746f2e Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 6 Dec 2020 12:16:56 +0200
Subject: [PATCH] Add PackageCollectorRequirement to allow multiple entries of
 the same package

---
 trains_agent/commands/worker.py             |  3 +-
 trains_agent/helper/package/priority_req.py | 35 +++++++++++++++++++++
 trains_agent/helper/package/requirements.py | 16 ++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/trains_agent/commands/worker.py b/trains_agent/commands/worker.py
index 3f0345e..92320cb 100644
--- a/trains_agent/commands/worker.py
+++ b/trains_agent/commands/worker.py
@@ -90,7 +90,7 @@ from trains_agent.helper.process import (
     get_docker_id,
     commit_docker, terminate_process,
 )
-from trains_agent.helper.package.priority_req import PriorityPackageRequirement
+from trains_agent.helper.package.priority_req import PriorityPackageRequirement, PackageCollectorRequirement
 from trains_agent.helper.repo import clone_repository_cached, RepoInfo, VCS
 from trains_agent.helper.resource_monitor import ResourceMonitor
 from trains_agent.helper.runtime_verification import check_runtime, print_uptime_properties
@@ -322,6 +322,7 @@ class Worker(ServiceCommandSection):
         PriorityPackageRequirement,
         PostRequirement,
         ExternalRequirements,
+        partial(PackageCollectorRequirement, collect_package=['trains']),
     )
 
     # poll queues every _polling_interval seconds
diff --git a/trains_agent/helper/package/priority_req.py b/trains_agent/helper/package/priority_req.py
index 8f1c037..7459e6e 100644
--- a/trains_agent/helper/package/priority_req.py
+++ b/trains_agent/helper/package/priority_req.py
@@ -38,3 +38,38 @@ class PriorityPackageRequirement(SimpleSubstitution):
             return Text('')
         PackageManager.out_of_scope_install_package(str(req))
         return Text(req)
+
+
+class PackageCollectorRequirement(SimpleSubstitution):
+    """
+    This RequirementSubstitution class will allow you to have multiple instances of the same
+    package, it will output the last one (by order) to be actually used.
+    """
+    name = tuple()
+
+    def __init__(self, session, collect_package):
+        super(PackageCollectorRequirement, self).__init__(session)
+        self._collect_packages = collect_package
+        self._last_req = None
+
+    def match(self, req):
+        # match package names
+        return req.name and (req.name.lower() in self.name or req.name.lower() in self._collect_packages)
+
+    def replace(self, req):
+        """
+        Replace a requirement
+        :raises: ValueError if version is pre-release
+        """
+        self._last_req = req.clone()
+        return ''
+
+    def post_scan_add_req(self):
+        """
+        Allows the RequirementSubstitution to add an extra line/requirements after
+        the initial requirements scan is completed.
+        Called only once per requirements.txt object
+        """
+        last_req = self._last_req
+        self._last_req = None
+        return last_req
diff --git a/trains_agent/helper/package/requirements.py b/trains_agent/helper/package/requirements.py
index ee981ca..c187ad6 100644
--- a/trains_agent/helper/package/requirements.py
+++ b/trains_agent/helper/package/requirements.py
@@ -338,6 +338,14 @@ class RequirementSubstitution(object):
         """
         pass
 
+    def post_scan_add_req(self):  # type: () -> Optional[MarkerRequirement]
+        """
+        Allows the RequirementSubstitution to add an extra line/requirements after
+        the initial requirements scan is completed.
+        Called only once per requirements.txt object
+        """
+        return None
+
     def post_install(self, session):
         pass
 
@@ -492,6 +500,14 @@ class RequirementsManager(object):
         )
         if not conda:
             result = map(self.translator.translate, result)
+
+        result = list(result)
+        # add post scan add requirements call back
+        for h in self.handlers:
+            req = h.post_scan_add_req()
+            if req:
+                result.append(req.tostr())
+
         return join_lines(result)
 
     def post_install(self, session):