Python best practice: How to handle function parameters that can be 2 types?

user_2771455 0
Viewed 86

I am building a Python Class that works with PySpark dataframes or Hive tables. The input data can either be a str for table_name or a DataFrame. What is the Python best practice to do this? I wanted to do this so that it is flexible depending on the use case. The only difference is on how the data should be passed to the class (either as a table or a Spark DataFrame), everything else would be the same.

See example below:

class DataPipeline:

    def __init__(self, data):
        if isinstance(data, str):
            self.df = spark.read.table(data)
        elif isinstance(data, DataFrame):
            self.df = data
        else:
            raise ValueError("some error")

    def process_data(self):
       # do something with the self.df here
pyspark python
2 Answers

Don't. Make __init__ expect a DataFrame, and make the caller responsible (perhaps via a class method) for converting a string to a DataFrame before DataPipeline.__init__ is called.

class DataPipeline:

    def __init__(self, data: DataFrame):
        self.df = df


    @classmethod
    def from_string(cls, data: str):
        return cls(spark.read.table(data))


p1 = DataPipeline(some_data_frame)
p2 = DataPipeline.from_string("...")  # DataPipline(spark.read.table("..."))

Ryan Singer (supposedly; for as often as it is quoted, I've yet to find the original source) once said

So much complexity in software comes from trying to make one thing do two things.

Here, you are trying to make one thing (DataPipeline.__init__) to two things (initialize a DataPipeline using a string, and initialize a DataPipeline with a DataFrame).

For the refactoring, I chose to make the simplest case (accepting a DataFrame) for __init__, with the more complex logic of first parsing a string into a DataFrame moved into a separate class method that calls __init__ with an appropriate data frame.

(Whether you want to make any explicit runtime checks for type validation to fail early, or just assume that the caller will accept the consequences of failing to pass a value of the correct type, is up to you.)

user_11268410

You have the right basic idea but the syntax is wrong -- use raise to raise an exception(I'd suggest TypeError for an invalid type rather than ValueError), and make sure you're setting self.df instead of the local var df:

class DataPipeline:

    def __init__(self, data):
        if isinstance(data, str):
            self.df = spark.read.table(data)
        elif isinstance(data, DataFrame):
            self.df = data
        else:
            raise TypeError("Data must be str or DataFrame")
user_37997590
Related