cleanup conda/pypi usage of _user_defined_attributes

Netflix · Dec 9, 2024 · ff1c433 · ff1c433
1 parent 3c5a3f2
commit ff1c433
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 28 deletions.
diff --git 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
@@ -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"],