typing_example_ds.py

import numpy as np
import pandas as pd
from sklearn import linear_model
from sklearn.model_selection import train_test_split

# In general for the way to do type hinting is very straight forward.
# simply don't instantiate the class to use it as a 'type'

# Here's a generic example:

class ReturnType:
    def __init__(self):
        pass

class ParameterType:
    def __init__(self):
        pass

def some_function(a: ParameterType) -> ReturnType:
    return a

# Now let's see how this works for pandas:
# all we need to do is not instantiate the object:

def create_df() -> pd.DataFrame:
    return pd.DataFrame(
    np.random.normal(0, 1000, size=1000).reshape(200, 5),
    columns=["A", "B", "C", "D", "E"])

# We can use this with anything:

def regress_data(model: linear_model.LinearRegression, data: pd.DataFrame) -> np.array:
    X = data[["A", "B", "C", "D"]]
    y = data["E"]
    X_train, X_test, y_train, y_test = train_test_split(X, y)
    model.fit(X_train, y_train)
    return model.predict(X_test)