pythonpycharmpython-typing

How do you annotate for custom classes?


Say that I have a class, Shard, and my return value for a method should be a list of Shard objects.

my_method(args) -> ??:

Pycharm kicks up an error when I imported List from typing and specified List[Shard] (the error is related to it not being able to find Shard).

Here is the method in question:

def _return_shards(self, sql: str = None) -> List[Shard]:
        """
        Queries the database to return a sequence of all the database clusters
        and shards, and then creates Shard objects with the information.

        @type conn_info: dict
        :param conn_info: Database information containing the host, port, user,
         database, and password.
         @type sql: str

        :return: A list of Shard items
        """
        logging.info('Creating shards.')
        cursor = self._return_cursor(self.canvas_conn)

        stock_query = '''
              SELECT
                database_server_id || 'dr' database_server_id
              , CASE
                  WHEN database_server_id LIKE 'cluster1' AND name IS NULL THEN 'canvas'
                  WHEN database_server_id LIKE 'cluster4' AND name IS NULL THEN 'canvas'
                  WHEN name IS NULL THEN database_server_id
                  ELSE name
                END AS database_name
                ,SUBSTRING(database_server_id FROM '\d+$')::INT cluster_id
                ,CASE
                  WHEN name IS NULL THEN SUBSTRING(database_server_id FROM '\d+$')::INT
                  ELSE SUBSTRING(name FROM '\d+$')::INT
                END AS shard_id
              FROM canvas.switchman_shards
              WHERE database_server_id LIKE 'cluster%'
              ORDER BY cluster_id, shard_id
              '''
        if sql:
            cursor.execute(sql)
        else:
            cursor.execute(stock_query)
        results = cursor.fetchall()
        self._set_reporter_total_shards(count=len(results))
        return [Shard(
            *result,
            job=self.job,
            manifest=self.manifest,
            reporter=self.reporter.create_slave(db=result[0], schema=result[1])
        ) for result in results]

The error is NameError: name 'Shard' is not defined even though the Shard class is defined in the same module.


Solution

  • Shard isn't defined yet, either because you're still in the middle of the class definition, or because Shard is defined below this class. Give the name as a string:

    def _return_shards(self, sql: str = None) -> List['Shard']: