star

venv/lib/python3.9/site-packages/mlflow/entities/param.py

@@ -0,0 +1,49 @@
+import sys
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import Param as ProtoParam
+
+
+class Param(_MlflowObject):
+    """
+    Parameter object.
+    """
+
+    def __init__(self, key, value):
+        if "pyspark.ml" in sys.modules:
+            import pyspark.ml.param
+
+            if isinstance(key, pyspark.ml.param.Param):
+                key = key.name
+                value = str(value)
+        self._key = key
+        self._value = value
+
+    @property
+    def key(self):
+        """String key corresponding to the parameter name."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the parameter."""
+        return self._value
+
+    def to_proto(self):
+        param = ProtoParam()
+        param.key = self.key
+        param.value = self.value
+        return param
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self._key == __o._key
+
+        return False
+
+    def __hash__(self):
+        return hash(self._key)