From ff1c43332d327d5efb488d9519d21eff8f85fbbc Mon Sep 17 00:00:00 2001
From: Sakari Ikonen <sakari.a.ikonen@gmail.com>
Date: Mon, 9 Dec 2024 17:25:46 +0200
Subject: [PATCH] cleanup conda/pypi usage of _user_defined_attributes

---
 metaflow/plugins/pypi/conda_decorator.py | 21 +++------------------
 metaflow/plugins/pypi/pypi_decorator.py  | 13 +++----------
 2 files changed, 6 insertions(+), 28 deletions(-)

diff --git a/metaflow/plugins/pypi/conda_decorator.py b/metaflow/plugins/pypi/conda_decorator.py
index f881f7f6b0..fc0e4d83a1 100644
--- a/metaflow/plugins/pypi/conda_decorator.py
+++ b/metaflow/plugins/pypi/conda_decorator.py
@@ -49,23 +49,9 @@ class CondaStepDecorator(StepDecorator):
     # CONDA_CHANNELS in their environment. For pinning specific packages to specific
     # conda channels, users can specify channel::package as the package name.
 
-    def __init__(self, attributes=None, statically_defined=False):
-        self._user_defined_attributes = (
-            attributes.copy() if attributes is not None else {}
-        )
-        super(CondaStepDecorator, self).__init__(attributes, statically_defined)
-
     def init(self):
         super(CondaStepDecorator, self).init()
 
-        # We have to go back and fixup _user_defined_attributes for potential
-        # config resolution
-        self._user_defined_attributes = {
-            k: v
-            for k, v in self.attributes.items()
-            if k in self._user_defined_attributes
-        }
-
         # Support legacy 'libraries=' attribute for the decorator.
         self.attributes["packages"] = {
             **self.attributes["libraries"],
@@ -94,10 +80,9 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 **super_attributes["packages"],
                 **self.attributes["packages"],
             }
-            self._user_defined_attributes = {
-                **self._user_defined_attributes,
-                **conda_base._user_defined_attributes,
-            }
+            self._user_defined_attributes = self._user_defined_attributes.union(
+                conda_base._user_defined_attributes
+            )
             self.attributes["python"] = (
                 self.attributes["python"] or super_attributes["python"]
             )
diff --git a/metaflow/plugins/pypi/pypi_decorator.py b/metaflow/plugins/pypi/pypi_decorator.py
index 1b16967ced..575531eb45 100644
--- a/metaflow/plugins/pypi/pypi_decorator.py
+++ b/metaflow/plugins/pypi/pypi_decorator.py
@@ -24,12 +24,6 @@ class PyPIStepDecorator(StepDecorator):
     name = "pypi"
     defaults = {"packages": {}, "python": None, "disabled": None}  # wheels
 
-    def __init__(self, attributes=None, statically_defined=False):
-        self._user_defined_attributes = (
-            attributes.copy() if attributes is not None else {}
-        )
-        super().__init__(attributes, statically_defined)
-
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # The init_environment hook for Environment creates the relevant virtual
         # environments. The step_init hook sets up the relevant state for that hook to
@@ -42,10 +36,9 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         if "pypi_base" in self.flow._flow_decorators:
             pypi_base = self.flow._flow_decorators["pypi_base"][0]
             super_attributes = pypi_base.attributes
-            self._user_defined_attributes = {
-                **self._user_defined_attributes,
-                **pypi_base._user_defined_attributes,
-            }
+            self._user_defined_attributes = self._user_defined_attributes.union(
+                pypi_base._user_defined_attributes
+            )
             self.attributes["packages"] = {
                 **super_attributes["packages"],
                 **self.attributes["packages"],