type hinting, class inheritance and metaclass with a cython lib

Question

I am using a cython lib called pygrib that doesn't have any associated type hinting.

I would like to extend its features by wrapping, for example the pygrib.open method in some classes with relative type hints.

pygrib.open returns an Iterator with pygrib.messages

the classes

from abc import ABCMeta, abstractmethod
from typing import Iterator, Type, List
import pygrib

class GrbMessage:
    def keys(self) -> List[str]: ...


class GrbMeta(Type, metaclass=ABCMeta):
    @abstractmethod
    def __iter__(self) -> Iterator[GrbMessage]: ...


class GrbOpen(pygrib.open, metaclass=GrbMeta):
    ...

usage

If I explicitly state by using the change_type function that the variable is GrbMeta the VSCode ide detects the type hinting.

def change_type(x)->GrbMeta:
    return x

def _unzip(path: str) -> str:
    os.system(f'gzip -d {path}')
    return path.strip('.gz')

if __name__ == '__main__':
    grbs = GrbOpen(_unzip(GRIB_FILE))
    typed_grbs = change_type(grbs)

    for grb in grbs:
        grb.keys()

    for grb in typed_grbs:
        grb.keys()

A possible solution is to just return self from the from __iter__

possible solution

class GrbOpen(pygrib.open):

    def __iter__(self) -> Iterator[GrbMessage]:
        return self

I'm just curious if this could be done better.

Why does GrbOpen(pygrib.open, metaclass=GrbMeta) not inherit the GrbMeta typehints

Answer 1

It's often the process of asking a question that brings me to a solution.

Creating the NewType seems to have achieved what I was trying to accomplish.

GribFileType = NewType("GribFile", GribFileType)

class GribFile(pygrib.open):

    def to_dataframe(self: GribFileType):
        return pd.DataFrame.from_records(dict(_full_message(grb)) for grb in self)


class Reader:
    """reader wrapper with typing support around the pygrib open function"""
    grib_file: GribFile

    def __init__(self, file_path: str):
        if file_path.endswith('.gz'):
            self.file_path = _unzip(file_path)
        else:
            self.file_path = file_path

    def __enter__(self) -> GribFileType:
        self.grib_file = GribFile(self.file_path)
        return self.grib_file

    @abstractmethod
    def __exit__(self, *args) -> None:
        self.grib_file.close()