python-3.xabstract-classmypyconcreteclass

Python mypy marks error when method parameter is type Union

I have these python classes:

class LocalWritable(typing.TypedDict):
    file_name: str


class GSheetWritable(typing.TypedDict):
    tab_name: str


class S3Writable(typing.TypedDict):
    data_name: str
    table_name: str


WriterMeta = typing.Union[GSheetWritable, S3Writable, LocalWritable]

class DataWriter(ABC):
    """Defines the interface for all data writers"""

    @abstractmethod
    def write(self, data: pd.DataFrame, meta: WriterMeta, versionize: bool):
        """This method performs the writing of 'data'.

        Every class implementing this method must implement its writing
        using 'connector'
        """
        pass

class GSheetOutputWriter(DataWriter):
    def write(self, data: pd.DataFrame, meta: WriterMeta, versionize: bool):
        data = data.replace({np.nan: 0, np.Inf: "Inf"})

        print("Writing '{}' table to gsheet.".format(meta["tab_name"]))
        if self.new:
            tab = self.connector.get_worksheet(self.target.url, "Sheet1")
            self.connector.rename_worksheet(tab, meta["tab_name"])
            self.new = False
        else:
            tab = self.connector.add_worksheet(
                self.target, meta["tab_name"], rows=1, cols=1
            )

        time.sleep(random.randint(30, 60))
        self.connector.update_worksheet(
            tab, [data.columns.values.tolist()] + data.values.tolist()
        )

The problem is with the method write() when linting with python mypy, because it marks this error:

cost_reporter\outputs\__init__.py:209: error: TypedDict "S3Writable" has no key "tab_name"
cost_reporter\outputs\__init__.py:209: note: Did you mean "table_name" or "data_name"?
cost_reporter\outputs\__init__.py:209: error: TypedDict "LocalWritable" has no key "tab_name"

What I am trying to do is to implement three concrete classes based on the abstract class DataWriter, and each one shall implement its own write() method and each one shall receive one of the datatypes of WriterMeta union. The problem I am having is that python mypy validates the code against the three datatypes instead of any of them.

How can I do that?

EDIT

If I change the type of parameter meta to GsheetWritable(that is one of the three types of the union and the one expected by this concrete class), mypy marks this error:

cost_reporter\outputs\__init__.py:202: error: Argument 2 of "write" is incompatible with supertype "DataWriter"; supertype defines the argument type as "Union[GSheetWritable, S3Writable, LocalWritable]"
cost_reporter\outputs\__init__.py:202: note: This violates the Liskov substitution principle
Solution

  • A Union works like unions in set theory. In other words, a Union consisting of multiple types is a type that supports only what's shared in common.

    In order to use attributes (or whatever) of a specific type, you need to hint to mypy that you're constraining an instance. You can do this by casting the Union to a specific type, asserting that your object is whatever specific type, and others. The documentation lists ways to narrow types.

    import typing
from abc import ABC, abstractmethod

class LocalWritable(typing.TypedDict):
    file_name: str


class GSheetWritable(typing.TypedDict):
    tab_name: str


class S3Writable(typing.TypedDict):
    data_name: str
    table_name: str


WriterMeta = typing.Union[GSheetWritable, S3Writable, LocalWritable]


class DataWriter(ABC):
    @abstractmethod
    def write(self, data: str, meta: WriterMeta):
        pass


class GSheetOutputWriter(DataWriter):
    def write(self, data: str, meta: WriterMeta):
        # LOOK HERE! The cast hints to mypy that meta is a GSheetWritable.
        meta_cast: GSheetWritable = typing.cast(GSheetWritable, meta)
        print("Writing '{}' table to gsheet.".format(meta_cast["tab_name"]))

    Further reading