Repository: dirty-cat/dirty_cat
Pull request #287 (2022-07-08 15:43:24)

Maintenance
This PR aims at improving the overall quality of the code and doc.

It has several purposes:
- Correct typos
- Reword unclear sentences
- Minor updates to the doc
- Some minor structural improvements, such as moving some functions to suiting modules
- Rename some variables for better readability
- Use modern language features for better readability and performance
- Simplify the code, while leaving functionalities intact (no bug-fixes)
- Make extensive use of type hinting, which serves two purposes:
  - Make the code easier to work with, especially when working with IDEs that support type hinting
  - Make the functions more efficient and less error-prone when using tools that enforce types, such as MyPy

In general, these are rather small modifications for which making unique PRs would be kind of overkill.


Patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index a5a3bf2..225b441 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -3,12 +3,12 @@ version: 2
 jobs:
   python3:
     docker:
-      - image: circleci/python:3.8.1
+      - image: cimg/python:3.8
     environment:
       - MINICONDA_PATH: ~/miniconda
       - CONDA_ENV_NAME: testenv
-      - PYTHON_VERSION: 3.8.1
-      - OMP_NUM_THREADS: 2 # Avoid overcommit with HistGradientBoosting
+      - PYTHON_VERSION: 3.8
+      - OMP_NUM_THREADS: 2 # Avoid over-commit with HistGradientBoosting
     steps:
       - checkout
       - run: ./build_tools/circle/checkout_merge_commit.sh
@@ -31,7 +31,7 @@ jobs:
 
   deploy:
     docker:
-      - image: circleci/python:3.8.1
+      - image: cimg/python:3.8
     steps:
       - checkout
       - run: ./build_tools/circle/checkout_merge_commit.sh
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
index d4f5dd0..936162a 100644
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -4,21 +4,19 @@ about: Something not working as described? This is the place.
 title: ''
 labels: 'bug'
 assignees: ''
-
 ---
 <!--Provide a brief description of the bug.-->
 
 
 <!--Please fill in the following information, to the best of your ability.-->
-dirty_cat version:
-
-### Expected behavior
 
+dirty_cat version:
 
-### Actual behavior
+Expected behavior:
 
+Observed behavior:
 
-### Steps and code to reproduce bug
+Steps and code to reproduce the bug:
 
 ```python
 
diff --git a/.github/ISSUE_TEMPLATE/documentation.md b/.github/ISSUE_TEMPLATE/documentation.md
index 97c1268..ed69f9c 100644
--- a/.github/ISSUE_TEMPLATE/documentation.md
+++ b/.github/ISSUE_TEMPLATE/documentation.md
@@ -5,7 +5,9 @@ title: ''
 labels: Documentation
 assignees: ''
 ---
-<!--Describe your proposed enhancement in detail.-->
+<!--Describe your proposed enhancement in detail-->
 
-<!--List any pages that would be impacted by the enhancement.-->
-### Affected pages
+<!--List pages that would be impacted by the enhancement-->
+
+Affected pages:
+- 
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
index f2b9dcd..c98bc88 100644
--- a/.github/ISSUE_TEMPLATE/feature_request.md
+++ b/.github/ISSUE_TEMPLATE/feature_request.md
@@ -6,14 +6,14 @@ labels: 'enhancement'
 assignees: ''
 
 ---
-<!--Provide a brief description what you would like changes and why.-->
+<!--Provide a brief description of what you'd like to change and why-->
 
 
-<!--Please fill in the following information, to the best of your ability.-->
-### Benefits to the change
+<!--Please fill in the following information, to the best of your ability-->
 
+Benefits to the change:
 
-### Pseudocode for the new behavior, if applicable
+Pseudocode for the new behavior (if applicable):
 
 ```python
 
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 4ab2ccd..31af2ae 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -13,3 +13,4 @@ jobs:
                   repo-token: ${{ secrets.GITHUB_TOKEN  }}
                   artifact-path: 0/doc/index.html
                   circleci-jobs: python3
+                  jobtitle: Check the rendered docs here!
diff --git a/.gitignore b/.gitignore
index dc835a9..a682db9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -58,4 +58,3 @@ doc/_build
 doc/auto_examples
 doc/generated
 .DS_Store
-
diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index 6c8ff28..0000000
--- a/.travis.yml
+++ /dev/null
@@ -1,28 +0,0 @@
-language: python
-
-install:
-  - git clone --depth 1 git://github.com/astropy/ci-helpers.git
-  - source ci-helpers/travis/setup_conda.sh
-  - source build_tools/travis/install.sh
-
-script:
-  - pytest --cov=./
-  - codecov
-
-matrix:
-  include:
-    - name: "Python 3.6 with minimum dependencies (no pandas and python-Levenshtein)"
-      python: "3.6"
-      env:
-        - CONDA_DEPENDENCIES='distance joblib numpy scipy requests pytest-cov codecov coverage'
-        - SCIKIT_LEARN_VERSION='stable'
-    - name: "Python 3.8 with pandas"
-      python: "3.8"
-      env:
-        - CONDA_DEPENDENCIES='python-Levenshtein distance joblib numpy scipy requests pytest-cov codecov coverage pandas'
-        - SCIKIT_LEARN_VERSION='stable'
-    - name: "Python 3.9 with scikit-learn nightly-build"
-      python: "3.9"
-      env:
-        - CONDA_DEPENDENCIES='python-Levenshtein distance joblib numpy scipy requests pytest-cov codecov coverage pandas'
-        - SCIKIT_LEARN_VERSION='nightly'
diff --git a/CHANGES.rst b/CHANGES.rst
index 397bad4..569decb 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -9,6 +9,7 @@ Major changes
     in the SuperVectorizer for datetime columns.
 * Support for Python 3.6 and 3.7 has been dropped. Python >= 3.8 is now required.
 * Bumped minimum dependencies:
+  - sklearn>=0.22
   - scipy>=1.4.0
   - numpy>=1.17.3
 
diff --git a/LICENSE.txt b/LICENSE.txt
index 0f469af..69aa265 100644
--- a/LICENSE.txt
+++ b/LICENSE.txt
@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2008-2016, The joblib developers.
+Copyright (c) 2018-2022, The dirty_cat developers.
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
diff --git a/README.rst b/README.rst
index e6c27a5..1f09038 100644
--- a/README.rst
+++ b/README.rst
@@ -37,7 +37,7 @@ dirty_cat requires:
 - Python (>= 3.8)
 - NumPy (>= 1.17.3)
 - SciPy (>= 1.4.0)
-- scikit-learn (>= 0.21.0)
+- scikit-learn (>= 0.22.0)
 - pandas (>= 1.1.5)
 
 User installation
diff --git a/RELEASE_PROCESS.md b/RELEASE_PROCESS.md
index 8f38a30..ef820ba 100644
--- a/RELEASE_PROCESS.md
+++ b/RELEASE_PROCESS.md
@@ -5,7 +5,7 @@
 This document is aimed toward established contributors the project.
 
 If you're a user looking to install dirty_cat, please refer to
-the instructions in the [README](README.rst) !
+the instructions in the [README](README.rst)!
 
 ### Process
 
@@ -41,4 +41,4 @@ here are the main steps and appropriate resources:
 - [Packaging and distributing software - Python Packaging User Guide](https://packaging.python.org/guides/distributing-packages-using-setuptools/)
 - [Publishing (Perfect) Python Packages on PyPi](https://youtu.be/GIF3LaRqgXo)
 - [Managing releases in a repository](https://docs.github.com/en/repositories/releasing-projects-on-github/managing-releases-in-a-repository)
-- [Should I deploy today ?](https://shouldideploy.today/)
+- [Should I deploy today?](https://shouldideploy.today/)
diff --git a/benchmarks/encoders_time_memory_bench.py b/benchmarks/encoders_time_memory_bench.py
new file mode 100644
index 0000000..18b48e9
--- /dev/null
+++ b/benchmarks/encoders_time_memory_bench.py
@@ -0,0 +1,85 @@
+# Inspired from :
+# https://github.com/scikit-learn/scikit-learn/blob/main/benchmarks/bench_text_vectorizers.py
+# plot run time and memory usage of the different encoders using the SuperVectorizer method
+# test if auto_cast=True plays an important role
+
+"""
+
+To run this benchmark, you will need,
+
+ * dirty_cat
+ * pandas
+ * memory_profiler
+
+"""
+import timeit
+import itertools
+
+import numpy as np
+import pandas as pd
+from memory_profiler import memory_usage
+
+from dirty_cat.datasets import fetch_traffic_violations
+
+from dirty_cat import SimilarityEncoder, GapEncoder, MinHashEncoder
+from dirty_cat import SuperVectorizer
+
+n_repeat = 3
+
+
+def run_encoders(Encoder, X, **params):
+    def f():
+        enc = Encoder(**params)
+        enc.fit_transform(X)
+
+    return f
+
+
+data = fetch_traffic_violations()
+
+X = data[1][:5000]
+y = data[2][:5000]
+
+print("=" * 80 + "\n#" + "    Dirty_cat encoders benchmark" + "\n" + "=" * 80 + "\n")
+print(f"Using a subset of the traffic violations dataset ({len(X)} observations).")
+print("This benchmarks runs in ~1 min ...")
+
+res = []
+
+for Encoder, (auto_cast, high_card_cat_transformer) in itertools.product(
+    [SuperVectorizer],
+    [
+        (True, SimilarityEncoder()),
+        (True, GapEncoder()),
+        (True, MinHashEncoder()),
+        (False, SimilarityEncoder()),
+        (False, GapEncoder()),
+        (False, MinHashEncoder()),
+    ],
+):
+    params = {
+        "auto_cast": auto_cast,
+        "high_card_cat_transformer": high_card_cat_transformer,
+    }
+    bench = {"encoder": Encoder}
+    bench.update(params)
+    dt = timeit.repeat(run_encoders(Encoder, X, **params), number=1, repeat=n_repeat)
+    bench["time"] = f"{np.mean(dt):.3f} (+-{np.std(dt):.3f})"
+
+    mem_usage = memory_usage(run_encoders(Encoder, X, **params))
+
+    bench["memory"] = f"{np.max(mem_usage):.1f}"
+
+    res.append(bench)
+
+df = pd.DataFrame(res).set_index(["auto_cast", "high_card_cat_transformer", "encoder"])
+
+print("\n========== Run time performance (sec) ===========\n")
+print(
+    "Computing the mean and the standard deviation "
+    f"of the run time over {n_repeat} runs...\n"
+)
+print(df["time"].unstack(level=-1))
+
+print("\n=============== Memory usage (MB) ===============\n")
+print(df["memory"].unstack(level=-1))
diff --git a/build_tools/travis/install.sh b/build_tools/travis/install.sh
deleted file mode 100644
index d11e1f9..0000000
--- a/build_tools/travis/install.sh
+++ /dev/null
@@ -1,8 +0,0 @@
-
-if [[ "$SCIKIT_LEARN_VERSION" == "stable" ]]; then
-  pip install scikit-learn
-
-elif [[ "$SCIKIT_LEARN_VERSION" == "nightly" ]]; then
-  # https://scikit-learn.org/stable/developers/advanced_installation.html#installing-nightly-builds
-  pip install --pre --extra-index https://pypi.anaconda.org/scipy-wheels-nightly/simple scikit-learn
-fi
diff --git a/codecov.yml b/codecov.yml
index 1ce784a..726fcaf 100644
--- a/codecov.yml
+++ b/codecov.yml
@@ -12,10 +12,10 @@ coverage:
     patch:
       default:
         # Be tolerant on slight code coverage diff on PRs to limit
-        # noisy red coverage status on github PRs.
+        # noisy red coverage status on GitHub PRs.
         # Note The coverage stats are still uploaded
         # to codecov so that PR reviewers can see uncovered lines
-        # in the github diff if they install the codecov browser
+        # in the GitHub diff if they install the codecov browser
         # extension:
         # https://github.com/codecov/browser-extension
         target: auto
diff --git a/dirty_cat/datasets/__init__.py b/dirty_cat/datasets/__init__.py
index bd83591..17d0d34 100644
--- a/dirty_cat/datasets/__init__.py
+++ b/dirty_cat/datasets/__init__.py
@@ -1,7 +1,3 @@
-"""
-Datasets module of dirty cat
-"""
-
 from .fetching import get_data_dir
 from .fetching import fetch_medical_charge
 from .fetching import fetch_midwest_survey
diff --git a/dirty_cat/datasets/fetching.py b/dirty_cat/datasets/fetching.py
index 1a818f6..ee1e83c 100644
--- a/dirty_cat/datasets/fetching.py
+++ b/dirty_cat/datasets/fetching.py
@@ -1,15 +1,8 @@
-# -*- coding: utf-8 -*-
-
 """
 Fetching functions to retrieve example datasets, using
 Scikit-Learn's ``fetch_openml()`` function.
 """
 
-
-# Author:
-# Lilian Boulard <lilian@boulard.fr>
-# https://github.com/LilianBoulard
-
 # Future notes:
 # - Watch out for ``fetch_openml()`` API modifications:
 # as of january 2021, the function is marked as experimental.
@@ -17,15 +10,14 @@ Scikit-Learn's ``fetch_openml()`` function.
 
 import gzip
 import json
-import sklearn
 import warnings
 import pandas as pd
 
 from pathlib import Path
 from dataclasses import dataclass
 from typing import Union, Dict, Any, List
+from sklearn.datasets import fetch_openml
 
-from dirty_cat.utils import Version
 from dirty_cat.datasets.utils import get_data_dir
 
 # Directory where the ``.gz`` files containing the
@@ -35,24 +27,24 @@ from dirty_cat.datasets.utils import get_data_dir
 # ``Experimental`` so the structure might change in future releases.
 # This path will be concatenated to the dirty_cat data directory,
 # available via the function ``get_data_dir()``.
-DETAILS_DIRECTORY = "openml/openml.org/api/v1/json/data/"
+DETAILS_DIRECTORY: str = "openml/openml.org/api/v1/json/data/"
 
-# Same as above ; for the datasets features location.
-FEATURES_DIRECTORY = "openml/openml.org/api/v1/json/data/features/"
+# Same as above; for the datasets features location.
+FEATURES_DIRECTORY: str = "openml/openml.org/api/v1/json/data/features/"
 
-# Same as above ; for the datasets data location.
-DATA_DIRECTORY = "openml/openml.org/data/v1/download/"
+# Same as above; for the datasets data location.
+DATA_DIRECTORY: str = "openml/openml.org/data/v1/download/"
 
 # The IDs of the datasets, from OpenML.
 # For each dataset, its URL is constructed as follows:
-openml_url = "https://www.openml.org/d/{ID}"
-ROAD_SAFETY_ID = 42803
-OPEN_PAYMENTS_ID = 42738
-MIDWEST_SURVEY_ID = 42805
-MEDICAL_CHARGE_ID = 42720
-EMPLOYEE_SALARIES_ID = 42125
-TRAFFIC_VIOLATIONS_ID = 42132
-DRUG_DIRECTORY_ID = 43044
+openml_url: str = "https://www.openml.org/d/{ID}"
+ROAD_SAFETY_ID: int = 42803
+OPEN_PAYMENTS_ID: int = 42738
+MIDWEST_SURVEY_ID: int = 42805
+MEDICAL_CHARGE_ID: int = 42720
+EMPLOYEE_SALARIES_ID: int = 42125
+TRAFFIC_VIOLATIONS_ID: int = 42132
+DRUG_DIRECTORY_ID: int = 43044
 
 
 @dataclass(unsafe_hash=True)
@@ -147,7 +139,7 @@ def fetch_openml_dataset(dataset_id: int,
     data_gz_path = data_directory / DATA_DIRECTORY / f'{file_id}.gz'
 
     if not data_gz_path.is_file():
-        # This is a double-check.
+        # double-check.
         # If the data file does not exist, download the dataset.
         _download_and_write_openml_dataset(dataset_id=dataset_id,
                                            data_directory=data_directory)
@@ -188,23 +180,14 @@ def _download_and_write_openml_dataset(dataset_id: int,
         If there is no Internet connection.
 
     """
-    from sklearn.datasets import fetch_openml
-
-    fetch_kwargs = {}
-    if Version(sklearn.__version__) >= Version('0.22'):
-        fetch_kwargs.update({'as_frame': True})
-
     # The ``fetch_openml()`` function returns a Scikit-Learn ``Bunch`` object,
     # which behaves just like a ``namedtuple``.
     # However, we do not want to save this data into memory:
     # we will read it from the disk later.
-    #
-    # Raises ``ValueError`` if the ID is incorrect (does not exist on OpenML)
-    # and ``urllib.error.URLError`` if there is no Internet connection.
     fetch_openml(
         data_id=dataset_id,
         data_home=str(data_directory),
-        **fetch_kwargs
+        as_frame=True,
     )
 
 
@@ -222,7 +205,7 @@ def _read_json_from_gz(compressed_dir_path: Path) -> dict:
     -------
     dict
         The information contained in the file,
-        converted from plain-text JSON.
+        parsed from plain-text JSON.
 
     """
     if not compressed_dir_path.is_file():
@@ -555,9 +538,9 @@ def fetch_traffic_violations(load_dataframe: bool = True
 
     Description of the dataset:
     > This dataset contains traffic violation information from all electronic
-    traffic violations issued in the Montgomery County, MD. Any information that can be used
-    to uniquely identify the vehicle, the vehicle owner or the officer issuing
-    the violation will not be published.
+    traffic violations issued in the Montgomery County, MD. Any information
+    that can be used to uniquely identify the vehicle, the vehicle owner or
+    the officer issuing the violation will not be published.
 
     Returns
     -------
diff --git a/dirty_cat/datasets/utils.py b/dirty_cat/datasets/utils.py
index 64be270..452d149 100644
--- a/dirty_cat/datasets/utils.py
+++ b/dirty_cat/datasets/utils.py
@@ -2,13 +2,15 @@ import os
 from pathlib import Path
 
 
-def get_data_dir(name=None) -> Path:
+def get_data_dir(name: str = None) -> Path:
     """
     Returns the directory in which dirty_cat looks for data.
 
-    This is typically useful for the end-user to check where the data is downloaded and stored.
+    This is typically useful for the end-user to check
+    where the data is downloaded and stored.
     """
-    # assuming we are in datasets.utils, this calls the module
+    # Note: we stick to os.path instead of pathlib.Path because
+    # it's easier to test, and the functionality is the same.
     module_path = Path(os.path.dirname(__file__)).resolve()
     data_dir = module_path / 'data'
     if name is not None:
diff --git a/dirty_cat/datetime_encoder.py b/dirty_cat/datetime_encoder.py
index b6b016b..c2e6a42 100644
--- a/dirty_cat/datetime_encoder.py
+++ b/dirty_cat/datetime_encoder.py
@@ -2,7 +2,7 @@ import numpy as np
 import pandas as pd
 
 from warnings import warn
-from typing import List
+from typing import List, Literal, Dict, Optional
 
 from sklearn import __version__ as sklearn_version
 from sklearn.base import BaseEstimator, TransformerMixin
@@ -10,42 +10,61 @@ from sklearn.utils.validation import check_is_fitted
 
 from dirty_cat.utils import check_input, Version
 
-# Some functions need aliases
-WORD_TO_ALIAS = {"year": "Y", "month": "M", "day": "D", "hour": "H", "minute": "min", "second": "S",
-                 "millisecond": "ms", "microsecond": "us", "nanosecond": "N"}
-TIME_LEVELS = ["year", "month", "day", "hour", "minute", "second", "millisecond", "microsecond", "nanosecond"]
+
+WORD_TO_ALIAS: Dict[str, str] = {
+    "year": "Y", "month": "M", "day": "D",
+    "hour": "H", "minute": "min", "second": "S",
+    "millisecond": "ms", "microsecond": "us", "nanosecond": "N"
+}
+TIME_LEVELS: List[str] = list(WORD_TO_ALIAS.keys())
+AcceptedTimeValues = Literal["year", "month", "day",
+                             "hour", "minute", "second",
+                             "millisecond", "microsecond", "nanosecond"]
 
 
 class DatetimeEncoder(BaseEstimator, TransformerMixin):
     """
-    This encoder transforms each datetime column into several numeric columns corresponding to temporal features,
-    e.g year, month, day...
-    Constant extracted features are dropped ; for instance, if the year is always the same in a feature, the extracted "year" column won't be added.
-    If the dates are timezone aware, all the features extracted will correspond to the provided timezone.
+    This encoder transforms each datetime column into several numeric columns
+    corresponding to temporal features, e.g year, month, day...
+    Constant extracted features are dropped; for instance, if the year is
+    always the same in a feature, the extracted "year" column won't be added.
+    If the dates are timezone aware, all the features extracted will correspond
+    to the provided timezone.
 
     Parameters
     ----------
-    extract_until : {"year", "month", "day", "hour", "minute", "second", "millisecond", "microsecond", "nanosecond"}, default="hour"
-        Extract up to this granularity. If all features have not been extracted, add the "total_time" feature, which contains
-        the time to epoch (in seconds).
-        For instance, if you specify "day", only "year", "month", "day" and "total_time" features will be created.
+    extract_until : AcceptedTimeValues, default="hour"
+        Extract up to this granularity.
+        If all features have not been extracted, add the "total_time" feature,
+        which contains the time to epoch (in seconds).
+        For instance, if you specify "day", only "year", "month", "day" and
+        "total_time" features will be created.
     add_day_of_the_week: bool, default=False
-        Add day of the week feature (if day is extracted). This is a numerical feature from 0 (Monday) to 6 (Sunday).
+        Add day of the week feature (if day is extracted).
+        This is a numerical feature from 0 (Monday) to 6 (Sunday).
 
     Attributes
     ----------
+    n_features_in_: int
+        Number of features in the data seen during fit.
     n_features_out_: int
         Number of features of the transformed data.
-    features_per_column_: Dict[int, List[str]]
+    features_per_column_: typing.Dict[int, typing.List[str]]
         Dictionary mapping the index of the original columns
         to the list of features extracted for each column.
-    col_names_: List[str]
-        List of the names of the features of the input data, if input data was a pandas DataFrame, otherwise None.
+    col_names_: typing.Optional[typing.List[str]]
+        List of the names of the features of the input data,
+        if input data was a pandas DataFrame, otherwise None.
     """
 
+    n_features_in_: int
+    n_features_out_: int
+    features_per_column_: Dict[int, List[str]]
+    col_names_: Optional[List[str]]
+
     def __init__(self,
-                 extract_until="hour",
-                 add_day_of_the_week=False):
+                 extract_until: AcceptedTimeValues = "hour",
+                 add_day_of_the_week: bool = False):
         self.extract_until = extract_until
         self.add_day_of_the_week = add_day_of_the_week
 
@@ -58,10 +77,12 @@ class DatetimeEncoder(BaseEstimator, TransformerMixin):
     def _validate_keywords(self):
         if self.extract_until not in TIME_LEVELS:
             raise ValueError(
-                f'"extract_until" should be one of {TIME_LEVELS}, got {self.extract_until}.'
+                f'"extract_until" should be one of {TIME_LEVELS}, '
+                f'got {self.extract_until}. '
             )
 
-    def _extract_from_date(self, date_series, feature):
+    @staticmethod
+    def _extract_from_date(date_series: pd.Series, feature: str):
         if feature == "year":
             return pd.DatetimeIndex(date_series).year.to_numpy()
         elif feature == "month":
@@ -86,36 +107,44 @@ class DatetimeEncoder(BaseEstimator, TransformerMixin):
             tz = pd.DatetimeIndex(date_series).tz
             # Compute the time in seconds from the epoch time UTC
             if tz is None:
-                return (pd.to_datetime(date_series) - pd.Timestamp("1970-01-01")) // pd.Timedelta('1s')
+                return (pd.to_datetime(date_series) -
+                        pd.Timestamp("1970-01-01")) // pd.Timedelta('1s')
             else:
-                return (pd.DatetimeIndex(date_series).tz_convert("utc") - pd.Timestamp("1970-01-01", tz="utc")) // pd.Timedelta('1s')
+                return (pd.DatetimeIndex(date_series).tz_convert("utc") -
+                        pd.Timestamp("1970-01-01", tz="utc")) // pd.Timedelta('1s')
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None) -> "DatetimeEncoder":
         """
-        Fit the DatetimeEncoder to X. In practice, just stores which extracted features
-        are not constant.
+        Fit the DatetimeEncoder to X. In practice, just stores which extracted
+        features are not constant.
 
         Parameters
         ----------
         X : array-like, shape (n_samples, n_features)
             Data where each column is a datetime feature.
+        y : None
+            Unused, only here for compatibility.
 
         Returns
         -------
-        self
+        DatetimeEncoder
             Fitted DatetimeEncoder instance.
+
         """
         self._validate_keywords()
-        # Columns to extract for each column, before taking into account constant columns
-        self._to_extract = TIME_LEVELS[:TIME_LEVELS.index(self.extract_until) + 1]
+        # Columns to extract for each column,
+        # before taking into account constant columns
+        self._to_extract = TIME_LEVELS[:TIME_LEVELS.index(
+            self.extract_until) + 1]
         if self.add_day_of_the_week:
             self._to_extract.append("dayofweek")
         if isinstance(X, pd.DataFrame):
-            self.col_names_ = X.columns
+            self.col_names_ = X.columns.to_list()
         else:
             self.col_names_ = None
         X = check_input(X)
-        self.features_per_column_ = {}  # Features to extract for each column, after removing constant features
+        # Features to extract for each column, after removing constant features
+        self.features_per_column_ = {}
         for i in range(X.shape[1]):
             self.features_per_column_[i] = []
         # Check which columns are constant
@@ -123,39 +152,48 @@ class DatetimeEncoder(BaseEstimator, TransformerMixin):
             for feature in self._to_extract:
                 if np.nanstd(self._extract_from_date(X[:, i], feature)) > 0:
                     self.features_per_column_[i].append(feature)
-            # If some date features have not been extracted, then add the "total_time" feature, which contains the full
-            # time to epoch
-            remainder = (pd.to_datetime(X[:, i]) - pd.to_datetime(pd.DatetimeIndex(X[:, i]).floor(
-                WORD_TO_ALIAS[self.extract_until]))).seconds.to_numpy()
+            # If some date features have not been extracted, then add the
+            # "total_time" feature, which contains the full time to epoch
+            remainder = (pd.to_datetime(X[:, i]) - pd.to_datetime(
+                pd.DatetimeIndex(X[:, i]).floor(
+                    WORD_TO_ALIAS[self.extract_until]))).seconds.to_numpy()
             if np.nanstd(remainder) > 0:
                 self.features_per_column_[i].append("total_time")
 
         self.n_features_in_ = X.shape[1]
-        self.n_features_out_ = len(np.concatenate(list(self.features_per_column_.values())))
+        self.n_features_out_ = len(np.concatenate(
+            list(self.features_per_column_.values())))
 
         return self
 
-    def transform(self, X, y=None):
-        """ Transform X by replacing each datetime column with corresponding numerical features.
+    def transform(self, X, y=None) -> np.array:
+        """
+        Transform X by replacing each datetime column with
+        corresponding numerical features.
 
         Parameters
         ----------
         X : array-like, shape (n_samples, n_features)
             The data to transform, where each column is a datetime feature.
+        y : None
+            Unused, only here for compatibility.
 
         Returns
         -------
-        array, shape (n_samples, n_features_out_)
+        np.array, shape (n_samples, n_features_out_)
             Transformed input.
         """
-        check_is_fitted(self, attributes=["n_features_in_", "n_features_out_", "features_per_column_"])
+        check_is_fitted(self, attributes=["n_features_in_", "n_features_out_",
+                                          "features_per_column_"])
         X = check_input(X)
         if X.shape[1] != self.n_features_in_:
             raise ValueError(
-                f"The number of features in the input data ({X.shape[1]}) does not match the number of features "
-                f"seen during fit ({self.n_features_in_})."
+                f"The number of features in the input data ({X.shape[1]}) "
+                f"does not match the number of features "
+                f"seen during fit ({self.n_features_in_}). "
             )
-        # Create a new array with the extracted features, choosing only features that weren't constant during fit
+        # Create a new array with the extracted features,
+        # choosing only features that weren't constant during fit
         X_ = np.empty((X.shape[0], self.n_features_out_), dtype=np.float64)
         idx = 0
         for i in range(X.shape[1]):
@@ -168,9 +206,9 @@ class DatetimeEncoder(BaseEstimator, TransformerMixin):
         """
         Returns clean feature names with format "<column_name>_<new_feature>"
         if the original data has column names, otherwise with format
-        "<column_index>_<new_feature>". new_feature is one of ["year", "month",
-        "day", "hour", "minute", "second", "millisecond", "microsecond",
-        "nanosecond", "dayofweek"]
+        "<column_index>_<new_feature>" where `<new_feature>` is one of
+        ["year", "month", "day", "hour", "minute", "second", "millisecond",
+        "microsecond", "nanosecond", "dayofweek"].
         """
         feature_names = []
         for i in self.features_per_column_.keys():
@@ -193,4 +231,3 @@ class DatetimeEncoder(BaseEstimator, TransformerMixin):
                 stacklevel=2,
             )
         return self.get_feature_names_out()
-
diff --git a/dirty_cat/fast_hash.py b/dirty_cat/fast_hash.py
index 5bd56fd..cfb2248 100644
--- a/dirty_cat/fast_hash.py
+++ b/dirty_cat/fast_hash.py
@@ -10,22 +10,23 @@ The principle is as follows:
      max-int overflow). The corresponding operation defines a random
      order in the interval [-maxint, maxint]
   3. Computing this dot product over a sliding window (to compute it for
-     every n-gram is a convolution
-  4. We can then take the min (or the max) of the resulting sliding
-  window
-
+     every n-gram is a convolution)
+  4. We can then take the min (or the max) of the resulting sliding window
 """
+
 import functools
 import numpy as np
 
 # Precompute to avoid the cost and
-# cast to int32 to speedup the min 
+# cast to int32 to speed up the min
 MININT32 = np.int32(-2 ** (32 - 1))
 MAXINT32 = np.int32(2 ** (32 - 1) - 1)
 
+
 @functools.lru_cache(maxsize=1024)
 def gen_atom(atom_len, seed=0):
-    """ Generate a random integer array (atom).
+    """
+    Generate a random integer array (atom).
 
     Parameters
     ----------
@@ -47,7 +48,8 @@ def gen_atom(atom_len, seed=0):
 
 
 def ngram_min_hash(string, ngram_range=(2, 4), seed=0, return_minmax=False):
-    """ Compute the min/max hash of the ngrams of the string.
+    """
+    Compute the min/max hash of the ngrams of the string.
 
     Parameters
     ----------
@@ -60,12 +62,12 @@ def ngram_min_hash(string, ngram_range=(2, 4), seed=0, return_minmax=False):
         Integer used to seed the hashing function.
     return_minmax : bool, default=False
         If True, returns both the minhash and maxhash of the string.
-        Else, only returns the minhash. 
+        Else, only returns the minhash.
+
     Returns
     -------
     int or tuple 
         The min_hash or (min_hash, max_hash) of the n-grams of the string.
-
     """
     # Create a numerical 1D array from the string
     array = np.frombuffer(string.encode(), dtype='int8', count=len(string))
diff --git a/dirty_cat/gap_encoder.py b/dirty_cat/gap_encoder.py
index 841ebb5..92e55e5 100644
--- a/dirty_cat/gap_encoder.py
+++ b/dirty_cat/gap_encoder.py
@@ -14,21 +14,28 @@ The principle is as follows:
        with the Kullback-Leibler divergence as loss, and a Gamma prior on H.
        We thus optimize H and W with the multiplicative update method.
 """
+
 import warnings
 import numpy as np
+import pandas as pd
+
 from scipy import sparse
+from numpy.random import RandomState
+from typing import Tuple, Optional, Union, Literal, Dict, List, Generator
+
+from sklearn.cluster import KMeans
+from sklearn.neighbors import NearestNeighbors
 from sklearn import __version__ as sklearn_version
+from sklearn.utils.fixes import _object_dtype_isnan
+from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils import check_random_state, gen_batches
 from sklearn.utils.extmath import row_norms, safe_sparse_dot
-from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.feature_extraction.text import CountVectorizer, HashingVectorizer
-from sklearn.cluster import KMeans
-from sklearn.neighbors import NearestNeighbors
-from sklearn.utils.fixes import _object_dtype_isnan
-import pandas as pd
+
 from .utils import check_input
 from dirty_cat.utils import Version
 
+
 if Version(sklearn_version) == Version('0.22'):
     with warnings.catch_warnings():
         warnings.simplefilter('ignore')
@@ -50,12 +57,28 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
 
     """See GapEncoder's docstring."""
 
-    def __init__(self, n_components=10, batch_size=128, gamma_shape_prior=1.1,
-                 gamma_scale_prior=1.0, rho=.95, rescale_rho=False,
-                 hashing=False, hashing_n_features=2**12, init='k-means++',
-                 tol=1e-4, min_iter=2, max_iter=5, ngram_range=(2, 4),
-                 analyzer='char', add_words=False, random_state=None,
-                 rescale_W=True, max_iter_e_step=20):
+    rho_: float
+    H_dict_: Dict[np.array, np.array]
+
+    def __init__(self,
+                 n_components: int = 10,
+                 batch_size: int = 128,
+                 gamma_shape_prior: float = 1.1,
+                 gamma_scale_prior: float = 1.0,
+                 rho: float = .95,
+                 rescale_rho: bool = False,
+                 hashing: bool = False,
+                 hashing_n_features: int = 2**12,
+                 init: Literal["k-means++", "random", "k-means"] = 'k-means++',
+                 tol: float = 1e-4,
+                 min_iter: int = 2,
+                 max_iter: int = 5,
+                 ngram_range: Tuple[int, int] = (2, 4),
+                 analyzer: Literal["word", "char", "char_wb"] = 'char',
+                 add_words: bool = False,
+                 random_state: Optional[Union[int, RandomState]] = None,
+                 rescale_W: bool = True,
+                 max_iter_e_step: int = 20):
 
         self.ngram_range = ngram_range
         self.n_components = n_components
@@ -76,7 +99,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         self.rescale_W = rescale_W
         self.max_iter_e_step = max_iter_e_step
 
-    def _init_vars(self, X):
+    def _init_vars(self, X) -> Tuple[np.array, np.array, np.array]:
         """
         Build the bag-of-n-grams representation V of X and initialize
         the topics W.
@@ -87,7 +110,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
                  analyzer=self.analyzer, ngram_range=self.ngram_range,
                  n_features=self.hashing_n_features,
                  norm=None, alternate_sign=False)
-            if self.add_words: # Init a word counts vectorizer if needed
+            if self.add_words:  # Init a word counts vectorizer if needed
                 self.word_count_ = HashingVectorizer(
                      analyzer='word',
                      n_features=self.hashing_n_features,
@@ -104,11 +127,11 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         # Build the n-grams counts matrix unq_V on unique elements of X
         unq_X, lookup = np.unique(X, return_inverse=True)
         unq_V = self.ngrams_count_.fit_transform(unq_X)
-        if self.add_words: # Add word counts to unq_V
+        if self.add_words:  # Add word counts to unq_V
             unq_V2 = self.word_count_.fit_transform(unq_X)
             unq_V = sparse.hstack((unq_V, unq_V2), format='csr')
 
-        if not self.hashing: # Build n-grams/word vocabulary
+        if not self.hashing:  # Build n-grams/word vocabulary
             if Version(sklearn_version) < Version('1.0'):
                 self.vocabulary = self.ngrams_count_.get_feature_names()
             else:
@@ -132,11 +155,11 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         # Update self.H_dict_ with unique input strings and their activations
         self.H_dict_.update(zip(unq_X, unq_H))
         if self.rescale_rho:
-            # Make update rate per iteration independant of the batch_size
+            # Make update rate per iteration independent of the batch_size
             self.rho_ = self.rho ** (self.batch_size / len(X))
         return unq_X, unq_V, lookup
 
-    def _get_H(self, X):
+    def _get_H(self, X: np.array) -> np.array:
         """
         Return the bag-of-n-grams representation of X.
         """
@@ -145,7 +168,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
             h_out[:] = self.H_dict_[x]
         return H_out
 
-    def _init_w(self, V, X):
+    def _init_w(self, V: np.array, X) -> Tuple[np.array, np.array, np.array]:
         """
         Initialize the topics W.
         If self.init='k-means++', we use the init method of
@@ -168,14 +191,15 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
                     x_squared_norms=row_norms(V, squared=True),
                     random_state=self.random_state,
                     n_local_trials=None)
-                W = W + .1 # To avoid restricting topics to few n-grams only
+                W = W + .1  # To avoid restricting topics to a few n-grams only
         elif self.init == 'random':
             W = self.random_state.gamma(
                 shape=self.gamma_shape_prior, scale=self.gamma_scale_prior,
                 size=(self.n_components, self.n_vocab))
         elif self.init == 'k-means':
-            prototypes = get_kmeans_prototypes(
-                X, self.n_components, analyzer=self.analyzer, random_state=self.random_state)
+            prototypes = get_kmeans_prototypes(X, self.n_components,
+                                               analyzer=self.analyzer,
+                                               random_state=self.random_state)
             W = self.ngrams_count_.transform(prototypes).A + .1
             if self.add_words:
                 W2 = self.word_count_.transform(prototypes).A + .1
@@ -197,14 +221,15 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
                     W2 = W2 + .1
                 W = np.concatenate((W, W2), axis=0)
         else:
-            raise AttributeError(
-                'Initialization method %s does not exist.' % self.init)
+            raise ValueError(
+                f'Initialization method {self.init!r} does not exist. '
+            )
         W /= W.sum(axis=1, keepdims=True)
         A = np.ones((self.n_components, self.n_vocab)) * 1e-10
         B = A.copy()
         return W, A, B
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None) -> "GapEncoderColumn":
         """
         Fit the GapEncoder on batches of X.
 
@@ -212,15 +237,18 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         ----------
         X : array-like, shape (n_samples, )
             The string data to fit the model on.
+        y : None
+            Unused, only here for compatibility.
         
         Returns
         -------
         self
+            Fitting GapEncoderColumn instance.
         """
         # Copy parameter rho
         self.rho_ = self.rho
         # Check if first item has str or np.str_ type
-        assert isinstance(X[0], str), "ERROR: Input data is not string."
+        assert isinstance(X[0], str), "Input data is not string. "
         # Make n-grams counts matrix unq_V
         unq_X, unq_V, lookup = self._init_vars(X)
         n_batch = (len(X) - 1) // self.batch_size + 1
@@ -234,7 +262,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
               lookup, n=self.batch_size)):
                 if i == n_batch-1:
                     W_last = self.W_.copy()
-                # Update the activations unq_H
+                # Update activations unq_H
                 unq_H[unq_idx] = _multiplicative_update_h(
                     unq_V[unq_idx], self.W_, unq_H[unq_idx],
                     epsilon=1e-3, max_iter=self.max_iter_e_step,
@@ -252,7 +280,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
                         self.W_ - W_last) / np.linalg.norm(W_last)
 
             if (W_change < self.tol) and (n_iter_ >= self.min_iter - 1):
-                break # Stop if the change in W is smaller than the tolerance
+                break  # Stop if the change in W is smaller than the tolerance
 
         # Update self.H_dict_ with the learned encoded vectors (activations)
         self.H_dict_.update(zip(unq_X, unq_H))
@@ -264,32 +292,33 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         Use `get_feature_names_out` instead.
         """
         warnings.warn(
-            "get_feature_names is deprecated in scikit-learn > 1.0. "
-            "use get_feature_names_out instead",
+            "Following the changes in scikit-learn 1.0, "
+            "get_feature_names is deprecated. "
+            "Use get_feature_names_out instead. ",
             DeprecationWarning,
             stacklevel=2,
-            )
-        return self.get_feature_names_out(n_labels=n_labels,
-                                          prefix=prefix)
+        )
+        return self.get_feature_names_out(n_labels=n_labels, prefix=prefix)
 
-    def get_feature_names_out(self, n_labels=3, prefix=''):
+    def get_feature_names_out(self, n_labels: int = 3, prefix: str = '',
+                              ) -> List[str]:
         """
         Returns the labels that best summarize the learned components/topics.
-        For each topic, labels with highest activations are selected.
-        
+        For each topic, labels with the highest activations are selected.
+
         Parameters
         ----------
-        
         n_labels : int, default=3
             The number of labels used to describe each topic.
-        
+        prefix : str, default=""
+            Used as a prefix for the categories.
+
         Returns
         -------
-        
-        topic_labels : list of strings
+        topic_labels : typing.List[str]
             The labels that best describe each topic.
-        
         """
+
         vectorizer = CountVectorizer()
         vectorizer.fit(list(self.H_dict_.keys()))
         if Version(sklearn_version) < Version('1.0'):
@@ -308,7 +337,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         topic_labels = [prefix + ', '.join(label) for label in topic_labels]
         return topic_labels
 
-    def score(self, X):
+    def score(self, X) -> float:
         """
         Returns the Kullback-Leibler divergence between the n-grams counts
         matrix V of X, and its non-negative factorization HW.
@@ -320,9 +349,10 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
 
         Returns
         -------
-        kl_divergence : float.
+        float.
             The Kullback-Leibler divergence.
         """
+
         # Build n-grams/word counts matrix
         unq_X, lookup = np.unique(X, return_inverse=True)
         unq_V = self.ngrams_count_.transform(unq_X)
@@ -347,21 +377,23 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
             'kullback-leibler', square_root=False)
         return kl_divergence
 
-    def partial_fit(self, X, y=None):
+    def partial_fit(self, X, y=None) -> "GapEncoderColumn":
         """
         Partial fit of the GapEncoder on X.
-        To be used in a online learning procedure where batches of data are
+        To be used in an online learning procedure where batches of data are
         coming one by one.
 
         Parameters
         ----------
         X : array-like, shape (n_samples, )
             The string data to fit the model on.
+        y : None
+            Unused, only here for compatibility.
         
         Returns
         -------
-        self
-        
+        GapEncoderColumn
+            The fitted GapEncoderColumn instance.
         """
         
         # Init H_dict_ with empty dict if it's the first call of partial_fit
@@ -371,9 +403,9 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         if not hasattr(self, 'rho_'):
             self.rho_ = self.rho
         # Check if first item has str or np.str_ type
-        assert isinstance(X[0], str), "ERROR: Input data is not string."
+        assert isinstance(X[0], str), "Input data is not string. "
         # Check if it is not the first batch
-        if hasattr(self, 'vocabulary'): # Update unq_X, unq_V with new batch
+        if hasattr(self, 'vocabulary'):  # Update unq_X, unq_V with new batch
             unq_X, lookup = np.unique(X, return_inverse=True)
             unq_V = self.ngrams_count_.transform(unq_X)
             if self.add_words:
@@ -397,7 +429,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
             unq_X, unq_V, lookup = self._init_vars(X)
 
         unq_H = self._get_H(unq_X)
-        # Update the activations unq_H
+        # Update unq_H, the activations
         unq_H = _multiplicative_update_h(
             unq_V, self.W_, unq_H,
             epsilon=1e-3, max_iter=self.max_iter_e_step,
@@ -412,7 +444,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         self.H_dict_.update(zip(unq_X, unq_H))
         return self
 
-    def _add_unseen_keys_to_H_dict(self, X):
+    def _add_unseen_keys_to_H_dict(self, X) -> None:
         """
         Add activations of unseen string categories from X to H_dict.
         """
@@ -427,7 +459,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
                 unseen_V, np.ones((unseen_V.shape[0], self.n_components)))
             self.H_dict_.update(zip(unseen_X, unseen_H))
 
-    def transform(self, X):
+    def transform(self, X) -> np.array:
         """
         Return the encoded vectors (activations) H of input strings in X.
         Given the learnt topics W, the activations H are tuned to fit V = HW.
@@ -443,7 +475,7 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
             Transformed input.
         """
         # Check if first item has str or np.str_ type
-        assert isinstance(X[0], str), "ERROR: Input data is not string."
+        assert isinstance(X[0], str), "Input data is not string. "
         unq_X = np.unique(X)
         # Build the n-grams counts matrix V for the string data to encode
         unq_V = self.ngrams_count_.transform(unq_X)
@@ -454,11 +486,11 @@ class GapEncoderColumn(BaseEstimator, TransformerMixin):
         self._add_unseen_keys_to_H_dict(unq_X)
         unq_H = self._get_H(unq_X)
         # Loop over batches
-        for slice in gen_batches(n=unq_H.shape[0],
-                                 batch_size=self.batch_size):
+        for slc in gen_batches(n=unq_H.shape[0],
+                               batch_size=self.batch_size):
             # Given the learnt topics W, optimize H to fit V = HW
-            unq_H[slice] = _multiplicative_update_h(
-                unq_V[slice], self.W_, unq_H[slice],
+            unq_H[slc] = _multiplicative_update_h(
+                unq_V[slc], self.W_, unq_H[slc],
                 epsilon=1e-3, max_iter=100,
                 rescale_W=self.rescale_W,
                 gamma_shape_prior=self.gamma_shape_prior,
@@ -479,79 +511,60 @@ class GapEncoder(BaseEstimator, TransformerMixin):
 
     Parameters
     ----------
-
     n_components : int, default=10
         Number of latent categories used to model string data.
-
     batch_size : int, default=128
         Number of samples per batch.
-
     gamma_shape_prior : float, default=1.1
         Shape parameter for the Gamma prior distribution.
-
     gamma_scale_prior : float, default=1.0
         Scale parameter for the Gamma prior distribution.
-
     rho : float, default=0.95
         Weight parameter for the update of the W matrix.
-
     rescale_rho : bool, default=False
         If true, use rho ** (batch_size / len(X)) instead of rho to obtain an
         update rate per iteration that is independent of the batch size.
-
     hashing : bool, default=False
         If true, HashingVectorizer is used instead of CountVectorizer.
         It has the advantage of being very low memory scalable to large
         datasets as there is no need to store a vocabulary dictionary in
         memory.
-
     hashing_n_features : int, default=2**12
         Number of features for the HashingVectorizer. Only relevant if
         hashing=True.
-
-    init : str, default='k-means++'
+    init : typing.Literal["k-means++", "random", "k-means"], default='k-means++'
         Initialization method of the W matrix.
         Options: {'k-means++', 'random', 'k-means'}.
         If init='k-means++', we use the init method of sklearn.cluster.KMeans.
         If init='random', topics are initialized with a Gamma distribution.
         If init='k-means', topics are initialized with a KMeans on the n-grams
         counts. This usually makes convergence faster but is a bit slower.
-
     tol : float, default=1e-4
         Tolerance for the convergence of the matrix W.
-
     min_iter : int, default=2
         Minimum number of iterations on the input data.
-
     max_iter : int, default=5
         Maximum number of iterations on the input data.
-
-    ngram_range : tuple, default=(2, 4)
+    ngram_range : typing.Tuple[int, int], default=(2, 4)
         The range of ngram length that will be used to build the
         bag-of-n-grams representation of the input data.
-
-    analyzer : str, default='char'.
+    analyzer : typing.Literal["word", "char", "char_wb"], default='char'.
         Analyzer parameter for the CountVectorizer/HashingVectorizer.
         Options: {‘word’, ‘char’, ‘char_wb’}, describing whether the matrix V
         to factorize should be made of word counts or character n-gram counts.
         Option ‘char_wb’ creates character n-grams only from text inside word
         boundaries; n-grams at the edges of words are padded with space.
-
     add_words : bool, default=False
         If true, add the words counts to the bag-of-n-grams representation
         of the input data.
-
-    random_state : int or None, default=None
+    random_state : typing.Optional[Union[int, RandomState]], default=None
         Pass an int for reproducible output across multiple function calls.
-
     rescale_W : bool, default=True
         If true, the weight matrix W is rescaled at each iteration
-        to have an l1 norm equal to 1 for each row.
-
+        to have a l1 norm equal to 1 for each row.
     max_iter_e_step : int, default=20
         Maximum number of iterations to adjust the activations h at each step.
-
-    handle_missing : 'error' or 'empty_impute' (default)
+    handle_missing : typing.Literal["error", "empty_impute"], default=empty_impute
         Whether to raise an error or impute with empty string '' if missing
         values (NaN) are present during fit (default is to impute).
         In the inverse transform, the missing category will be denoted as None.
@@ -559,6 +572,9 @@ class GapEncoder(BaseEstimator, TransformerMixin):
 
     Attributes
     ----------
+    rho_: float
+    fitted_models_: typing.List[GapEncoderColumn]
+    column_names_: typing.List[str]
 
     References
     ----------
@@ -568,12 +584,30 @@ class GapEncoder(BaseEstimator, TransformerMixin):
     
     """
 
-    def __init__(self, n_components=10, batch_size=128, gamma_shape_prior=1.1,
-                 gamma_scale_prior=1.0, rho=.95, rescale_rho=False,
-                 hashing=False, hashing_n_features=2**12, init='k-means++',
-                 tol=1e-4, min_iter=2, max_iter=5, ngram_range=(2, 4),
-                 analyzer='char', add_words=False, random_state=None,
-                 rescale_W=True, max_iter_e_step=20, handle_missing='zero_impute'):
+    rho_: float
+    fitted_models_: List[GapEncoderColumn]
+    column_names_: List[str]
+
+    def __init__(self,
+                 n_components: int = 10,
+                 batch_size: int = 128,
+                 gamma_shape_prior: float = 1.1,
+                 gamma_scale_prior: float = 1.0,
+                 rho: float = .95,
+                 rescale_rho: bool = False,
+                 hashing: bool = False,
+                 hashing_n_features: int = 2**12,
+                 init: Literal["k-means++", "random", "k-means"] = 'k-means++',
+                 tol: float = 1e-4,
+                 min_iter: int = 2,
+                 max_iter: int = 5,
+                 ngram_range: Tuple[int, int] = (2, 4),
+                 analyzer: Literal["word", "char", "char_wb"] = 'char',
+                 add_words: bool = False,
+                 random_state: Optional[Union[int, RandomState]] = None,
+                 rescale_W: bool = True,
+                 max_iter_e_step: int = 20,
+                 handle_missing: Literal["error", "empty_impute"] = 'zero_impute'):
 
         self.ngram_range = ngram_range
         self.n_components = n_components
@@ -595,7 +629,7 @@ class GapEncoder(BaseEstimator, TransformerMixin):
         self.max_iter_e_step = max_iter_e_step
         self.handle_missing = handle_missing
 
-    def _more_tags(self):
+    def _more_tags(self) -> Dict[str, List[str]]:
         """
         Used internally by sklearn to ease the estimator checks.
         """
@@ -629,21 +663,21 @@ class GapEncoder(BaseEstimator, TransformerMixin):
         """
         if self.handle_missing not in ['error', 'zero_impute']:
             raise ValueError(
-                "handle_missing should be either 'error' or "
-                f"'zero_impute', got {self.handle_missing!r}"
+                f"handle_missing should be either 'error' or "
+                f"'zero_impute', got {self.handle_missing!r}. "
             )
 
         missing_mask = _object_dtype_isnan(X)
 
         if missing_mask.any():
             if self.handle_missing == 'error':
-                raise ValueError('Input data contains missing values.')
+                raise ValueError('Input data contains missing values. ')
             elif self.handle_missing == 'zero_impute':
                 X[missing_mask] = ''
 
         return X
             
-    def fit(self, X, y=None):
+    def fit(self, X, y=None) -> "GapEncoder":
         """
         Fit the GapEncoder on batches of X.
 
@@ -651,12 +685,15 @@ class GapEncoder(BaseEstimator, TransformerMixin):
         ----------
         X : array-like, shape (n_samples, n_features)
             The string data to fit the model on.
+        y : None
+            Unused, only here for compatibility.
         
         Returns
         -------
-        self
-        
+        GapEncoder
+            Fitted GapEncoder instance.
         """
+
         # Copy parameter rho
         self.rho_ = self.rho
         # If X is a dataframe, store its column names
@@ -671,14 +708,14 @@ class GapEncoder(BaseEstimator, TransformerMixin):
             self.fitted_models_.append(col_enc.fit(X[:, k]))
         return self
 
-    def transform(self, X):
+    def transform(self, X) -> np.array:
         """
         Return the encoded vectors (activations) H of input strings in X.
         Given the learnt topics W, the activations H are tuned to fit V = HW.
         When X has several columns, they are encoded separately and
         then concatenated.
         
-        Remark: calling transform mutliple times in a row on the same
+        Remark: calling transform multiple times in a row on the same
         input X can give slightly different encodings. This is expected
         due to a caching mechanism to speed things up.        
 
@@ -692,6 +729,7 @@ class GapEncoder(BaseEstimator, TransformerMixin):
         H : 2-d array, shape (n_samples, n_topics * n_features)
             Transformed input.
         """
+
         # Check input data shape
         X = check_input(X)
         X = self._handle_missing(X)
@@ -701,22 +739,25 @@ class GapEncoder(BaseEstimator, TransformerMixin):
         X_enc = np.hstack(X_enc)
         return X_enc
 
-    def partial_fit(self, X, y=None):
+    def partial_fit(self, X, y=None) -> "GapEncoder":
         """
         Partial fit of the GapEncoder on X.
-        To be used in a online learning procedure where batches of data are
+        To be used in an online learning procedure where batches of data are
         coming one by one.
 
         Parameters
         ----------
         X : array-like, shape (n_samples, n_features)
             The string data to fit the model on.
-        
+        y : None
+            Unused, only here for compatibility.
+
         Returns
         -------
-        self
-        
+        GapEncoder
+            Fitted GapEncoder instance.
         """
+
         # If X is a dataframe, store its column names
         if isinstance(X, pd.DataFrame):
             self.column_names_ = list(X.columns)
@@ -733,53 +774,56 @@ class GapEncoder(BaseEstimator, TransformerMixin):
             self.fitted_models_[k].partial_fit(X[:, k])
         return self
 
-    def get_feature_names_out(self, col_names=None, n_labels=3):
+    def get_feature_names_out(self,
+                              col_names: Optional[Union[Literal["auto"],
+                                                        List[str]]] = None,
+                              n_labels: int = 3):
         """
         Returns the labels that best summarize the learned components/topics.
-        For each topic, labels with highest activations are selected.
-        
+        For each topic, labels with the highest activations are selected.
+
         Parameters
         ----------
-        
-        col_names : {None, list or str}, default=None
+        col_names : typing.Optional[typing.Union[typing.Literal["auto"], typing.List[str]]], default=None
             The column names to be added as prefixes before the labels.
             If col_names == None, no prefixes are used.
             If col_names == 'auto', column names are automatically defined:
-                - if the input data was a dataframe, its column names are used
-                - otherwise, 'col1', ..., 'colN' are used as prefixes
+                - if the input data was a dataframe, its column names are used,
+                - otherwise, 'col1', ..., 'colN' are used as prefixes.
             Prefixes can be manually set by passing a list  for col_names.
-            
         n_labels : int, default=3
             The number of labels used to describe each topic.
-        
+
         Returns
         -------
-        
         topic_labels : list of strings
             The labels that best describe each topic.
-        
         """
         assert hasattr(self, 'fitted_models_'), (
             'ERROR: GapEncoder must be fitted first.')
         # Generate prefixes
         if isinstance(col_names, str) and col_names == 'auto':
-            if hasattr(self, 'column_names_'): # Use column names
-                prefixes = [s + ': ' for s in self.column_names_]
-            else: # Use 'col1: ', ... 'colN: ' as prefixes
-                prefixes = [f'col{k}: ' for k in range(len(self.fitted_models_))]
+            if hasattr(self, 'column_names_'):  # Use column names
+                prefixes = [
+                    '%s: ' % col for col in self.column_names_
+                ]
+            else:  # Use 'col1: ', ... 'colN: ' as prefixes
+                prefixes = [
+                    'col%d: ' % i for i in range(len(self.fitted_models_))
+                ]
         elif col_names is None:  # Empty prefixes
             prefixes = [''] * len(self.fitted_models_)
         else:
-            prefixes = [s + ': ' for s in col_names]
+            prefixes = ['%s: ' % col for col in col_names]
         labels = list()
         for k, enc in enumerate(self.fitted_models_):
             col_labels = enc.get_feature_names_out(n_labels, prefixes[k])
             labels.extend(col_labels)
         return labels
     
-    def get_feature_names(
-        self, input_features=None, col_names=None, n_labels=3
-    ):
+    def get_feature_names(self, input_features=None,
+                          col_names: List[str] = None, n_labels: int = 3
+                          ) -> List[str]:
         """
         Ensures compatibility with sklearn < 1.0.
         Use `get_feature_names_out` instead.
@@ -788,14 +832,13 @@ class GapEncoder(BaseEstimator, TransformerMixin):
             warnings.warn(
                 "Following the changes in scikit-learn 1.0, "
                 "get_feature_names is deprecated. "
-                "Use get_feature_names_out instead.",
+                "Use get_feature_names_out instead. ",
                 DeprecationWarning,
                 stacklevel=2,
             )
         return self.get_feature_names_out(col_names, n_labels)
-        
 
-    def score(self, X):
+    def score(self, X) -> float:
         """
         Returns the sum over the columns of X of the Kullback-Leibler
         divergence between the n-grams counts matrix V of X, and its
@@ -808,26 +851,30 @@ class GapEncoder(BaseEstimator, TransformerMixin):
 
         Returns
         -------
-        kl_divergence : float.
+        float.
             The Kullback-Leibler divergence.
         """
         X = check_input(X)
         kl_divergence = 0
         for k in range(X.shape[1]):
-            kl_divergence += self.fitted_models_[k].score(X[:,k])
+            kl_divergence += self.fitted_models_[k].score(X[:, k])
         return kl_divergence
-        
-def _rescale_W(W, A):
+
+
+def _rescale_W(W: np.array, A: np.array) -> None:
     """
     Rescale the topics W to have a L1-norm equal to 1.
+    Note that they are modified in-place.
     """
     s = W.sum(axis=1, keepdims=True)
     W /= s
     A /= s
-    return
 
 
-def _multiplicative_update_w(Vt, W, A, B, Ht, rescale_W, rho):
+def _multiplicative_update_w(Vt: np.array, W: np.array, A: np.array,
+                             B: np.array, Ht: np.array, rescale_W: bool,
+                             rho: float,
+                             ) -> Tuple[np.array, np.array, np.array]:
     """
     Multiplicative update step for the topics W.
     """
@@ -841,7 +888,7 @@ def _multiplicative_update_w(Vt, W, A, B, Ht, rescale_W, rho):
     return W, A, B
 
 
-def _rescale_h(V, H):
+def _rescale_h(V: np.array, H: np.array) -> np.array:
     """
     Rescale the activations H.
     """
@@ -851,9 +898,11 @@ def _rescale_h(V, H):
     return H
 
 
-def _multiplicative_update_h(Vt, W, Ht, epsilon=1e-3, max_iter=10,
-                             rescale_W=False,
-                             gamma_shape_prior=1.1, gamma_scale_prior=1.):
+def _multiplicative_update_h(Vt: np.array, W: np.array, Ht: np.array,
+                             epsilon: float = 1e-3, max_iter: int = 10,
+                             rescale_W: bool = False,
+                             gamma_shape_prior: float = 1.1,
+                             gamma_scale_prior: float = 1.):
     """
     Multiplicative update step for the activations H.
     """
@@ -882,18 +931,26 @@ def _multiplicative_update_h(Vt, W, Ht, epsilon=1e-3, max_iter=10,
     return Ht
 
 
-def batch_lookup(lookup, n=1):
-    """ Make batches of the lookup array. """
+def batch_lookup(lookup: np.array, n: int = 1,
+                 ) -> Generator[Tuple[np.array, np.array], None, None]:
+    """
+    Make batches of the lookup array.
+    """
     len_iter = len(lookup)
     for idx in range(0, len_iter, n):
         indices = lookup[slice(idx, min(idx + n, len_iter))]
         unq_indices = np.unique(indices)
-        yield (unq_indices, indices)
+        yield unq_indices, indices
 
 
-def get_kmeans_prototypes(X, n_prototypes, analyzer='char', hashing_dim=128,
-                          ngram_range=(2, 4), sparse=False,
-                          sample_weight=None, random_state=None):
+def get_kmeans_prototypes(X,
+                          n_prototypes: int,
+                          analyzer: Literal["word", "char", "char_wb"] = 'char',
+                          hashing_dim: int = 128,
+                          ngram_range: Tuple[int, int] = (2, 4),
+                          sparse: bool = False,
+                          sample_weight=None,
+                          random_state: Optional[Union[int, RandomState]] = None):
     """
     Computes prototypes based on:
       - dimensionality reduction (via hashing n-grams)
diff --git a/dirty_cat/minhash_encoder.py b/dirty_cat/minhash_encoder.py
index e2c32ca..d41aea7 100644
--- a/dirty_cat/minhash_encoder.py
+++ b/dirty_cat/minhash_encoder.py
@@ -17,11 +17,15 @@ morphological similarities between strings.
 """
 
 import numpy as np
+
+from typing import Tuple, Literal, Dict, List
+
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils import murmurhash3_32
 
 from .fast_hash import ngram_min_hash
 from .utils import LRUDict, check_input
+from .string_distances import get_unique_ngrams
 
 
 class MinHashEncoder(BaseEstimator, TransformerMixin):
@@ -35,18 +39,24 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
     n_components : int, default=30
         The number of dimension of encoded strings. Numbers around 300 tend to
         lead to good prediction performance, but with more computational cost.
-    ngram_range : tuple (min_n, max_n), default=(2, 4)
+    ngram_range : typing.Tuple[int, int], default=(2, 4)
         The lower and upper boundary of the range of n-values for different
         n-grams to be extracted. All values of n such that min_n <= n <= max_n.
         will be used.
-    hashing : str {'fast', 'murmur'}, default=fast
+    hashing : typing.Literal["fast", "murmur"], default=fast
         Hashing function. fast is faster but
         might have some concern with its entropy.
     minmax_hash : bool, default=False
         if True, return min hash and max hash concatenated.
-    handle_missing : 'error' or 'zero_impute' (default)
+    handle_missing : typing.Literal["error", "zero_impute"], default=zero_impute
         Whether to raise an error or encode missing values (NaN) with
         vectors filled with zeros.
+
+    Attributes
+    ----------
+    hash_dict_ : LRUDict
+        Computed hashes.
+
     References
     ----------
     For a detailed description of the method, see
@@ -54,65 +64,53 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
     <https://hal.inria.fr/hal-02171256v4>`_ by Cerda, Varoquaux (2019).
 
     """
-    _capacity = 2 ** 10
+    hash_dict_: LRUDict
+
+    _capacity: int = 2 ** 10
 
-    def __init__(self, n_components=30, ngram_range=(2, 4),
-                 hashing='fast', minmax_hash=False,
-                 handle_missing='zero_impute'):
+    def __init__(self,
+                 n_components: int = 30,
+                 ngram_range: Tuple[int, int] = (2, 4),
+                 hashing: Literal["fast", "murmur"] = 'fast',
+                 minmax_hash: bool = False,
+                 handle_missing: Literal["error", "zero_impute"] = 'zero_impute'):
         self.ngram_range = ngram_range
         self.n_components = n_components
         self.hashing = hashing
         self.minmax_hash = minmax_hash
         self.handle_missing = handle_missing
 
-    def _more_tags(self):
+    def _more_tags(self) -> Dict[str, List[str]]:
         """
         Used internally by sklearn to ease the estimator checks.
         """
         return {"X_types": ["categorical"]}
 
-    def get_unique_ngrams(self, string, ngram_range):
-        """ Return the set of unique n-grams of a string.
-        Parameters
-        ----------
-        string : str
-            The string to split in n-grams.
-        ngram_range : tuple (min_n, max_n)
-        The lower and upper boundary of the range of n-values for different
-        n-grams to be extracted. All values of n such that min_n <= n <= max_n.
-        Returns
-        -------
-        set
-            The set of unique n-grams of the string.
+    def minhash(self, string: str, n_components: int,
+                ngram_range: Tuple[int, int]) -> np.array:
         """
-        spaces = ' '  # * (n // 2 + n % 2)
-        string = spaces + " ".join(string.lower().split()) + spaces
-        ngram_set = set()
-        for n in range(ngram_range[0], ngram_range[1] + 1):
-            string_list = [string[i:] for i in range(n)]
-            ngram_set |= set(zip(*string_list))
-        return ngram_set
-
-    def minhash(self, string, n_components, ngram_range):
-        """ Encode a string using murmur hashing function.
+        Encode a string using murmur hashing function.
+
         Parameters
         ----------
         string : str
             The string to encode.
         n_components : int
             The number of dimension of encoded string.
-        ngram_range : tuple (min_n, max_n)
-        The lower and upper boundary of the range of n-values for different
-        n-grams to be extracted. All values of n such that min_n <= n <= max_n.
+        ngram_range : typing.Tuple[int, int]
+            The lower and upper boundaries of the range of n-values for
+            different n-grams to be extracted.
+            All values of n such that min_n <= n <= max_n.
+
         Returns
         -------
         array, shape (n_components, )
             The encoded string.
         """
         min_hashes = np.ones(n_components) * np.infty
-        grams = self.get_unique_ngrams(string, self.ngram_range)
+        grams = get_unique_ngrams(string, self.ngram_range)
         if len(grams) == 0:
-            grams = self.get_unique_ngrams(' Na ', self.ngram_range)
+            grams = get_unique_ngrams(' Na ', self.ngram_range)
         for gram in grams:
             hash_array = np.array([
                 murmurhash3_32(''.join(gram), seed=d, positive=True)
@@ -120,17 +118,19 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
             min_hashes = np.minimum(min_hashes, hash_array)
         return min_hashes / (2 ** 32 - 1)
 
-    def get_fast_hash(self, string):
+    def get_fast_hash(self, string: str) -> np.array:
         """
         Encode a string with fast hashing function.
         fast hashing supports both min_hash and minmax_hash encoding.
+
         Parameters
         ----------
         string : str
             The string to encode.
+
         Returns
         -------
-        array, shape (n_components, )
+        np.array of shape (n_components, )
             The encoded string, using specified encoding scheme.
         """
         if self.minmax_hash:
@@ -141,30 +141,36 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
             return np.array([ngram_min_hash(string, self.ngram_range, seed)
                              for seed in range(self.n_components)])
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None) -> "MinHashEncoder":
         """
         Fit the MinHashEncoder to X. In practice, just initializes a dictionary
         to store encodings to speed up computation.
+
         Parameters
         ----------
         X : array-like, shape (n_samples, ) or (n_samples, 1)
             The string data to encode.
+        y : None
+            Unused, only here for compatibility.
 
         Returns
         -------
-        self
+        MinHashEncoder
             The fitted MinHashEncoder instance.
         """
         self.count = 0
-        self.hash_dict = LRUDict(capacity=self._capacity)
+        self.hash_dict_ = LRUDict(capacity=self._capacity)
         return self
 
-    def transform(self, X):
-        """ Transform X using specified encoding scheme.
+    def transform(self, X) -> np.array:
+        """
+        Transform X using specified encoding scheme.
+
         Parameters
         ----------
         X : array-like, shape (n_samples, ) or (n_samples, 1)
             The string data to encode.
+
         Returns
         -------
         array, shape (n_samples, n_components)
@@ -173,16 +179,12 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
         X = check_input(X)
         if self.minmax_hash:
             assert self.n_components % 2 == 0, \
-                "n_components should be even when minmax_hash=True"
+                "n_components should be even when minmax_hash=True. "
         if self.hashing == 'murmur':
-            assert not (self.minmax_hash), \
-                "minmax_hash not implemented with murmur"
-        if self.handle_missing not in ['error', 'zero_impute']:
-            template = ("handle_missing should be either 'error' or "
-                        "'zero_impute', got %s")
-            raise ValueError(template % self.handle_missing)
-
-        # TODO Parallel run here
+            assert not self.minmax_hash, \
+                'minmax_hash is not implemented with hashing="murmur". '
+
+        # TODO: Parallelize
         is_nan_idx = False
 
         if self.hashing == 'fast':
@@ -193,10 +195,10 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
                 for i, x in enumerate(X_in):
                     if isinstance(x, float):  # true if x is a missing value
                         is_nan_idx = True
-                    elif x not in self.hash_dict:
-                        X_out[i, k * self.n_components:counter] = self.hash_dict[x] = self.get_fast_hash(x)
+                    elif x not in self.hash_dict_:
+                        X_out[i, k * self.n_components:counter] = self.hash_dict_[x] = self.get_fast_hash(x)
                     else:
-                        X_out[i, k * self.n_components:counter] = self.hash_dict[x]
+                        X_out[i, k * self.n_components:counter] = self.hash_dict_[x]
                 counter += self.n_components
         elif self.hashing == 'murmur':
             X_out = np.zeros((len(X[:]), self.n_components * X.shape[1]))
@@ -206,22 +208,33 @@ class MinHashEncoder(BaseEstimator, TransformerMixin):
                 for i, x in enumerate(X_in):
                     if isinstance(x, float):
                         is_nan_idx = True
-                    elif x not in self.hash_dict:
-                        X_out[i, k * self.n_components:counter] = self.hash_dict[x] = self.minhash(
+                    elif x not in self.hash_dict_:
+                        X_out[i, k * self.n_components:counter] = self.hash_dict_[x] = self.minhash(
                             x,
                             n_components=self.n_components,
                             ngram_range=self.ngram_range
                         )
                     else:
-                        X_out[i, k * self.n_components:counter] = self.hash_dict[x]
+                        X_out[i, k * self.n_components:counter] = self.hash_dict_[x]
                 counter += self.n_components
         else:
-            raise ValueError("hashing function must be 'fast' or"
-                             "'murmur', got '{}'"
-                             "".format(self.hashing))
-
-        if self.handle_missing == 'error' and is_nan_idx:
-            msg = ("Found missing values in input data; set "
-                   "handle_missing='zero_impute' to encode with missing values")
-            raise ValueError(msg)
+            raise ValueError(
+                f'Got hashing={self.hashing}, '
+                f'but expected any of {{"fast", "murmur"}}. '
+            )
+
+        if self.handle_missing == 'error':
+            if is_nan_idx:
+                raise ValueError(
+                    "Found missing values in input data; set "
+                    "handle_missing='zero_impute' "
+                    "to encode with missing values. "
+                )
+        elif self.handle_missing == 'zero_impute':
+            pass
+        else:
+            raise ValueError(
+                f'Got handle_missing={self.handle_missing}, but expected '
+                f'any of {{"error", "zero_impute"}}. '
+            )
         return X_out
diff --git a/dirty_cat/similarity_encoder.py b/dirty_cat/similarity_encoder.py
index f5efe1d..d5eb138 100644
--- a/dirty_cat/similarity_encoder.py
+++ b/dirty_cat/similarity_encoder.py
@@ -12,11 +12,15 @@ The principle is as follows:
        be automatically sampled from the input data (most frequent categories,
        KMeans) or provided by the user.
 """
-import warnings
 
+import warnings
 import numpy as np
-from joblib import Parallel, delayed
+
 from scipy import sparse
+from numpy.random import RandomState
+from joblib import Parallel, delayed
+from typing import List, Tuple, Literal, Optional, Union
+
 import sklearn
 from sklearn.cluster import KMeans
 from sklearn.feature_extraction.text import CountVectorizer, HashingVectorizer
@@ -30,50 +34,60 @@ from .string_distances import get_ngram_count, preprocess
 
 
 def _ngram_similarity_one_sample_inplace(
-        x_count_vector, vocabulary_count_matrix, str_x,
-        vocabulary_ngram_counts, se_dict, unq_X, i, ngram_range):
-    """Update inplace a dict of similarities between a string and a vocabulary
-
+        x_count_vector: np.array,
+        vocabulary_count_matrix: np.array,
+        str_x: str,
+        vocabulary_ngram_counts: np.array,
+        se_dict: dict,
+        unq_X: np.array,
+        i: int,
+        ngram_range: Tuple[int, int]) -> None:
+    """
+    Update inplace a dict of similarities between a string and a vocabulary
 
     Parameters
     ----------
-    x_count_vector: np.array
-        count vector of the sample based on the ngrams of the vocabulary
-    vocabulary_count_matrix: np.array
-        count vector of the vocabulary based on its ngrams
+    x_count_vector : np.array
+        Count vector of the sample based on the ngrams of the vocabulary
+    vocabulary_count_matrix : np.array
+        Count vector of the vocabulary based on its ngrams
     str_x: str
-        the actual sample string
-    vocabulary_ngram_counts: np.array
-        number of ngrams for each unique element of the vocabulary
-    se_dict: dict
-        dictionary containing the similarities for each x in unq_X
-    unq_X: np.array
-        the arrayes of all unique samples
-    i: str
-        the index of x_count_vector in the csr count matrix
-    ngram_range: tuple
-
+        The actual sample string
+    vocabulary_ngram_counts : np.array
+        Number of ngrams for each unique element of the vocabulary
+    se_dict : dict
+        Dictionary containing the similarities for each x in unq_X
+    unq_X : np.array
+        The arrays of all unique samples
+    i : str
+        The index of x_count_vector in the csr count matrix
+    ngram_range : tuple
+        n-grams to use for the decomposition, where `n_min <= n <= n_max`.
     """
     nonzero_idx = x_count_vector.indices
     nonzero_vals = x_count_vector.data
 
-    samegrams = np.asarray(
+    same_grams = np.asarray(
         (vocabulary_count_matrix[:, nonzero_idx].minimum(nonzero_vals)).sum(
             axis=1))
 
-    allgrams = get_ngram_count(
-        str_x, ngram_range) + vocabulary_ngram_counts - samegrams
-    similarity = np.divide(samegrams, allgrams, out=np.zeros_like(samegrams),
-                           where=allgrams != 0)
+    all_grams = get_ngram_count(
+        str_x, ngram_range) + vocabulary_ngram_counts - same_grams
+    similarity = np.divide(same_grams, all_grams,
+                           out=np.zeros_like(same_grams),
+                           where=all_grams != 0)
     se_dict[unq_X[i]] = similarity.reshape(-1)
 
 
-def ngram_similarity(X, cats, ngram_range, hashing_dim, dtype=np.float64):
+def ngram_similarity(X,
+                     cats: List[str],
+                     ngram_range: Tuple[int, int],
+                     hashing_dim: int,
+                     dtype: type = np.float64) -> np.array:
     """
     Similarity encoding for dirty categorical variables:
-        Given to arrays of strings, returns the
-        similarity encoding matrix of size
-        len(X) x len(cats)
+    Given two arrays of strings, returns the similarity encoding matrix
+    of size len(X) x len(cats)
 
     ngram_sim(s_i, s_j) =
         ||min(ci, cj)||_1 / (||ci||_1 + ||cj||_1 - ||min(ci, cj)||_1)
@@ -103,10 +117,10 @@ def ngram_similarity(X, cats, ngram_range, hashing_dim, dtype=np.float64):
 
     for i, x in enumerate(count_X):
         _, nonzero_idx, nonzero_vals = sparse.find(x)
-        samegrams = np.asarray(
+        same_grams = np.asarray(
             (count_cats[:, nonzero_idx].minimum(nonzero_vals)).sum(axis=1))
-        allgrams = x.sum() + sum_cats - samegrams
-        similarity = np.divide(samegrams, allgrams)
+        all_grams = x.sum() + sum_cats - same_grams
+        similarity = np.divide(same_grams, all_grams)
         SE_dict[unq_X[i]] = similarity.reshape(-1)
     # We don't need the counts anymore, delete them to save memory
     del count_cats, count_X
@@ -118,7 +132,7 @@ def ngram_similarity(X, cats, ngram_range, hashing_dim, dtype=np.float64):
     return np.nan_to_num(out, copy=False)
 
 
-def get_prototype_frequencies(prototypes):
+def get_prototype_frequencies(prototypes: np.array) -> np.array:
     """
     Computes the frequencies of the values contained in prototypes
     Reverse sorts the array by the frequency
@@ -129,9 +143,14 @@ def get_prototype_frequencies(prototypes):
     return uniques[sorted_indexes], counts[sorted_indexes]
 
 
-def get_kmeans_prototypes(X, n_prototypes, hashing_dim=128,
-                          ngram_range=(3, 3), sparse=False, sample_weight=None,
-                          random_state=None):
+def get_kmeans_prototypes(X,
+                          n_prototypes: int,
+                          hashing_dim: int = 128,
+                          ngram_range: Tuple[int, int] = (3, 3),
+                          sparse: bool = False,
+                          sample_weight=None,
+                          random_state: Optional[Union[int, RandomState]] = None
+                          ) -> np.array:
     """
     Computes prototypes based on:
       - dimensionality reduction (via hashing n-grams)
@@ -153,12 +172,13 @@ def get_kmeans_prototypes(X, n_prototypes, hashing_dim=128,
     indexes_prototypes = np.unique(neighbors.kneighbors(centers, 1)[-1])
     if indexes_prototypes.shape[0] < n_prototypes:
         warnings.warn('Final number of unique prototypes is lower than ' +
-                      'n_prototypes (expected)')
+                      'n_prototypes (expected). ')
     return np.sort(X[indexes_prototypes])
 
 
 class SimilarityEncoder(OneHotEncoder):
-    """Encode string categorical features as a numeric array.
+    """
+    Encode string categorical features as a numeric array.
 
     The input to this transformer should be an array-like of
     strings.
@@ -173,14 +193,10 @@ class SimilarityEncoder(OneHotEncoder):
         Deprecated in dirty_cat 0.3, will be removed in 0.5.
         Was used to specify the type of pairwise string similarity to use.
         Since 0.3, only the ngram similarity is supported.
-
     ngram_range : tuple (min_n, max_n), default=(2, 4)
         The range of values for the n_gram similarity.
-
-    categories : 'auto', 'k-means', 'most_frequent' or a list of lists/arrays
-    of values.
+    categories : typing.Union[typing.Literal["auto", "k-means", "most_frequent"], typing.List[typing.List[str]]]
         Categories (unique values) per feature:
-
         - 'auto' : Determine categories automatically from the training data.
         - list : ``categories[i]`` holds the categories expected in the i-th
           column. The passed categories must be sorted and should not mix
@@ -189,13 +205,12 @@ class SimilarityEncoder(OneHotEncoder):
            categorical variable
         - 'k-means' : Computes the K nearest neighbors of K-mean centroids
            in order to choose the prototype categories
-
         The categories used can be found in the ``categories_`` attribute.
     dtype : number type, default np.float64
         Desired dtype of output.
     handle_unknown : 'error' or 'ignore' (default)
-        Whether to raise an error or ignore if a unknown categorical feature is
-        present during transform (default is to ignore). When this parameter
+        Whether to raise an error or ignore if an unknown categorical feature
+        is present during transform (default is to ignore). When this parameter
         is set to 'ignore' and an unknown category is encountered during
         transform, the resulting encoded columns for this feature
         will be all zeros. In the inverse transform, an unknown category
@@ -210,22 +225,22 @@ class SimilarityEncoder(OneHotEncoder):
     hashing_dim : int type or None.
         If None, the base vectorizer is CountVectorizer, else it's set to
         HashingVectorizer with a number of features equal to `hashing_dim`.
-    n_prototypes: number of prototype we want to use.
+    n_prototypes : number of prototype we want to use.
         Useful when `most_frequent` or `k-means` is used.
-        Must be a positive non null integer.
-    random_state: either an int used as a seed, a RandomState instance or None.
+        Must be a positive non-null integer.
+    random_state : either an int used as a seed, a RandomState instance or None.
         Useful when `k-means` strategy is used.
-    n_jobs: int, optional
+    n_jobs : int, optional
         maximum number of processes used to compute similarity matrices. Used
         only if ``fast=True`` in ``SimilarityEncoder.transform``
 
     Attributes
     ----------
-    categories_ : list of arrays
+    categories_ : typing.List[np.array]
         The categories of each feature determined during fitting
         (in order corresponding with output of ``transform``).
-    _infrequent_enabled: bool, default=False
-        Avoid taking into account the existance of infrequent categories.
+    _infrequent_enabled : bool, default=False
+        Avoid taking into account the existence of infrequent categories.
 
     References
     ----------
@@ -235,13 +250,28 @@ class SimilarityEncoder(OneHotEncoder):
     <https://hal.inria.fr/hal-01806175>`_ by Cerda, Varoquaux, Kegl. 2018
     (accepted for publication at: Machine Learning journal, Springer).
 
-
     """
 
-    def __init__(self, similarity=None, ngram_range=(2, 4),
-                 categories='auto', dtype=np.float64,
-                 handle_unknown='ignore', handle_missing='', hashing_dim=None,
-                 n_prototypes=None, random_state=None, n_jobs=None):
+    categories_: List[np.array]
+    n_features_in_: int
+    random_state_: Union[int, RandomState]
+    drop_idx_: np.array
+    vectorizers_: List[CountVectorizer]
+    vocabulary_count_matrices_: List[np.array]
+    vocabulary_ngram_counts_: List[List[int]]
+    _infrequent_enabled: bool
+
+    def __init__(self,
+                 similarity: str = 'ngram',
+                 ngram_range: Tuple[int, int] = (2, 4),
+                 categories: Union[Literal["auto", "k-means", "most_frequent"], List[List[str]]] = 'auto',
+                 dtype: type = np.float64,
+                 handle_unknown: Literal["error", "ignore"] = 'ignore',
+                 handle_missing: Literal["error", ""] = '',
+                 hashing_dim: Optional[int] = None,
+                 n_prototypes: Optional[int] = None,
+                 random_state: Optional[Union[int, RandomState]] = None,
+                 n_jobs: Optional[int] = None):
         super().__init__()
         self.categories = categories
         self.dtype = dtype
@@ -265,28 +295,32 @@ class SimilarityEncoder(OneHotEncoder):
 
         if not isinstance(categories, list):
             assert categories in [None, 'auto', 'k-means', 'most_frequent']
-        if categories in ['k-means', 'most_frequent'] and (n_prototypes is None
-                                                           or n_prototypes == 0
-                                                           ):
+        if categories in ['k-means', 'most_frequent'] \
+                and (n_prototypes is None or n_prototypes == 0):
             raise ValueError(
-                'n_prototypes expected None or a positive non null integer')
+                'n_prototypes expected None or a positive non null integer. ')
         if categories == 'auto' and n_prototypes is not None:
             warnings.warn(
-                'n_prototypes parameter ignored with category type \'auto\'')
+                'n_prototypes parameter ignored with category type "auto". ')
+
+    def get_most_frequent(self, prototypes: List[str]) -> np.array:
+        """
+        Get the most frequent category prototypes.
 
-    def get_most_frequent(self, prototypes):
-        """ Get the most frequent category prototypes
         Parameters
         ----------
-        prototypes : the list of values for a category variable
+        prototypes : typing.List[str]
+            The list of values for a category variable.
+
         Returns
         -------
-        The n_prototypes most frequent values for a category variable
+        np.array
+            The n_prototypes most frequent values for a category variable.
         """
         values, _ = get_prototype_frequencies(prototypes)
         return values[:self.n_prototypes]
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None) -> "SimilarityEncoder":
         """
         Fit the SimilarityEncoder to X.
 
@@ -294,21 +328,27 @@ class SimilarityEncoder(OneHotEncoder):
         ----------
         X : array-like, shape [n_samples, n_features]
             The data to determine the categories of each feature.
+        y : None
+            Unused, only here for compatibility.
+
         Returns
         -------
-        self
+        SimilarityEncoder
+            The fitted SimilarityEncoder instance.
         """
 
         if self.handle_missing not in ['error', '']:
-            template = ("handle_missing should be either 'error' or "
-                        "'', got %s")
-            raise ValueError(template % self.handle_missing)
+            raise ValueError(
+                f"Got handle_missing={self.handle_missing}, but expected "
+                f"any of {{'error', ''}}. "
+            )
         if hasattr(X, 'iloc') and X.isna().values.any():
             if self.handle_missing == 'error':
-                msg = ("Found missing values in input data; set "
-                       "handle_missing='' to encode with missing values")
-                raise ValueError(msg)
-            if self.handle_missing != 'error':
+                raise ValueError(
+                    "Found missing values in input data; set "
+                    "handle_missing='' to encode with missing values. "
+                )
+            else:
                 X = X.fillna(self.handle_missing)
         elif not hasattr(X, 'dtype') and isinstance(X, list):
             X = np.asarray(X, dtype=object)
@@ -317,31 +357,35 @@ class SimilarityEncoder(OneHotEncoder):
             mask = _object_dtype_isnan(X)
             if X.dtype.kind == 'O' and mask.any():
                 if self.handle_missing == 'error':
-                    msg = ("Found missing values in input data; set "
-                           "handle_missing='' to encode with missing values")
-                    raise ValueError(msg)
-                if self.handle_missing != 'error':
+                    raise ValueError(
+                        "Found missing values in input data; set "
+                        "handle_missing='' to encode with missing values. "
+                    )
+                else:
                     X[mask] = self.handle_missing
 
         Xlist, n_samples, n_features = self._check_X(X)
         self.n_features_in_ = n_features
 
         if self.handle_unknown not in ['error', 'ignore']:
-            template = ("handle_unknown should be either 'error' or "
-                        "'ignore', got %s")
-            raise ValueError(template % self.handle_unknown)
+            raise ValueError(
+                f"Got handle_unknown={self.handle_unknown!r}, but expected "
+                f"any of {{'error', 'ignore'}}. "
+            )
 
         if ((self.hashing_dim is not None) and
                 (not isinstance(self.hashing_dim, int))):
-            raise ValueError("value '%r' was specified for hashing_dim, "
-                             "which has invalid type, expected None or "
-                             "int." % self.hashing_dim)
+            raise ValueError(
+                f"Got hashing_dim={self.hashing_dim!r}, which has an invalid "
+                f"type ({type(self.hashing_dim)}), expected None or int. "
+            )
 
         if self.categories not in ['auto', 'most_frequent', 'k-means']:
             for cats in self.categories:
                 if not np.all(np.sort(cats) == np.array(cats)):
-                    raise ValueError("Unsorted categories are not yet "
-                                     "supported")
+                    raise ValueError(
+                        "Unsorted categories are not yet supported. "
+                    )
 
         self.categories_ = list()
         self.random_state_ = check_random_state(self.random_state)
@@ -355,17 +399,20 @@ class SimilarityEncoder(OneHotEncoder):
             elif self.categories == 'k-means':
                 uniques, count = np.unique(Xi, return_counts=True)
                 self.categories_.append(
-                    get_kmeans_prototypes(uniques, self.n_prototypes,
-                                          sample_weight=count,
-                                          random_state=self.random_state_))
+                    get_kmeans_prototypes(
+                        uniques, self.n_prototypes, sample_weight=count,
+                        random_state=self.random_state_,
+                    )
+                )
             else:
                 if self.handle_unknown == 'error':
                     valid_mask = np.in1d(Xi, self.categories[i])
                     if not np.all(valid_mask):
                         diff = np.unique(Xi[~valid_mask])
-                        msg = ("Found unknown categories {0} in column {1}"
-                               " during fit".format(diff, i))
-                        raise ValueError(msg)
+                        raise ValueError(
+                            f"Found unknown categories {diff} in column {i} "
+                            f"during fit. "
+                        )
                 self.categories_.append(np.array(self.categories[i],
                                                  dtype=object))
 
@@ -384,18 +431,20 @@ class SimilarityEncoder(OneHotEncoder):
             # equivalency between the user input and the categories_
             # attribute of the SimilarityEncoder, while being compliant
             # with the CountVectorizer preprocessing steps.
-            preprocessed_categories = np.array(list(map(
-                preprocess, self.categories_[i])), dtype=object)
+            categories = self.categories_[i]
 
-            vocabulary_count_matrix = vectorizer.fit_transform(
-                preprocessed_categories)
+            self.vectorizers_.append(vectorizer)
 
-            vocabulary_ngram_count = list(map(lambda x: get_ngram_count(
-                preprocess(x), self.ngram_range), self.categories_[i]))
+            self.vocabulary_count_matrices_.append(
+                vectorizer.fit_transform([
+                    preprocess(category) for category in categories
+                ])
+            )
 
-            self.vectorizers_.append(vectorizer)
-            self.vocabulary_count_matrices_.append(vocabulary_count_matrix)
-            self.vocabulary_ngram_counts_.append(vocabulary_ngram_count)
+            self.vocabulary_ngram_counts_.append([
+                get_ngram_count(preprocess(category), self.ngram_range)
+                for category in categories
+            ])
 
         self.drop_idx_ = self._compute_drop_idx()
         if Version(sklearn.__version__) >= Version('1.1.0'):
@@ -403,7 +452,7 @@ class SimilarityEncoder(OneHotEncoder):
 
         return self
 
-    def transform(self, X, fast=True):
+    def transform(self, X, fast: bool = True) -> np.array:
         """
         Transform X using specified encoding scheme.
 
@@ -411,19 +460,21 @@ class SimilarityEncoder(OneHotEncoder):
         ----------
         X : array-like, shape [n_samples, n_features]
             The data to encode.
+        fast : bool
+            Whether to use the fast computation of ngrams.
 
         Returns
         -------
         X_new : 2-d array, shape [n_samples, n_features_new]
             Transformed input.
-
         """
 
         if hasattr(X, 'iloc') and X.isna().values.any():
             if self.handle_missing == 'error':
-                msg = ("Found missing values in input data; set "
-                       "handle_missing='' to encode with missing values")
-                raise ValueError(msg)
+                raise ValueError(
+                    "Found missing values in input data; set "
+                    "handle_missing='' to encode with missing values. "
+                )
             if self.handle_missing != 'error':
                 X = X.fillna(self.handle_missing)
         elif not hasattr(X, 'dtype') and isinstance(X, list):
@@ -433,10 +484,11 @@ class SimilarityEncoder(OneHotEncoder):
             mask = _object_dtype_isnan(X)
             if X.dtype.kind == 'O' and mask.any():
                 if self.handle_missing == 'error':
-                    msg = ("Found missing values in input data; set "
-                           "handle_missing='' to encode with missing values")
-                    raise ValueError(msg)
-                if self.handle_missing != 'error':
+                    raise ValueError(
+                        "Found missing values in input data; set "
+                        "handle_missing='' to encode with missing values. "
+                    )
+                else:
                     X[mask] = self.handle_missing
 
         Xlist, n_samples, n_features = self._check_X(X)
@@ -448,50 +500,51 @@ class SimilarityEncoder(OneHotEncoder):
             if not np.all(valid_mask):
                 if self.handle_unknown == 'error':
                     diff = np.unique(X[~valid_mask, i])
-                    msg = ("Found unknown categories {0} in column {1}"
-                           " during transform".format(diff, i))
-                    raise ValueError(msg)
+                    raise ValueError(
+                        f"Found unknown categories {diff} in column {i} "
+                        f"during fit. "
+                    )
 
         min_n, max_n = self.ngram_range
 
         total_length = sum(len(x) for x in self.categories_)
         out = np.empty((len(X), total_length), dtype=self.dtype)
         last = 0
-        for j, cats in enumerate(self.categories_):
+        for j, categories in enumerate(self.categories_):
             if fast:
                 encoded_Xj = self._ngram_similarity_fast(Xlist[j], j)
             else:
                 encoded_Xj = ngram_similarity(
-                    Xlist[j], cats, ngram_range=(min_n, max_n),
+                    Xlist[j], categories, ngram_range=(min_n, max_n),
                     hashing_dim=self.hashing_dim, dtype=np.float32)
 
-            out[:, last:last + len(cats)] = encoded_Xj
-            last += len(cats)
+            out[:, last:last + len(categories)] = encoded_Xj
+            last += len(categories)
         return out
 
-    def _ngram_similarity_fast(self, X, col_idx):
+    def _ngram_similarity_fast(self, X: Union[list, np.array], col_idx: int,
+                               ) -> np.array:
         """
         Fast computation of ngram similarity.
 
-
         SimilarityEncoder.transform uses the count vectors of the vocabulary in
         its computations. In ngram_similarity, these count vectors have to be
         re-computed each time, which can slow down the execution. In this
         method, the count vectors are recovered from the
         ``vocabulary_count_matrices`` attribute of the SimilarityEncoder,
         speeding up the execution.
+
         Parameters
         ----------
-        X: np.array, list
-            observations being transformed.
+        X: typing.Union[list, np.array]
+            Observations being transformed.
         col_idx: int
-            the column index of X in the original feature matrix.
+            The column index of X in the original feature matrix.
         """
-        min_n, max_n = self.ngram_range
         vectorizer = self.vectorizers_[col_idx]
 
         unq_X = np.unique(X)
-        unq_X_ = np.array(list(map(preprocess, unq_X)))
+        unq_X_ = np.array([preprocess(x) for x in unq_X])
 
         X_count_matrix = vectorizer.transform(unq_X_)
         vocabulary_count_matrix = self.vocabulary_count_matrices_[col_idx]
@@ -500,45 +553,50 @@ class SimilarityEncoder(OneHotEncoder):
 
         se_dict = {}
 
-        Parallel(n_jobs=self.n_jobs, backend='threading')(delayed(
-            _ngram_similarity_one_sample_inplace)(
-            X_count_vector, vocabulary_count_matrix, x_str,
-            vocabulary_ngram_count, se_dict, unq_X, i, self.ngram_range) for
-            X_count_vector, x_str, i in zip(
-                X_count_matrix, unq_X_, range(len(unq_X))))
+        Parallel(n_jobs=self.n_jobs, backend='threading')(
+            delayed(
+                _ngram_similarity_one_sample_inplace)(
+                    X_count_vector, vocabulary_count_matrix, x_str,
+                    vocabulary_ngram_count, se_dict, unq_X, i, self.ngram_range
+            )
+            for X_count_vector, x_str, i
+            in zip(X_count_matrix, unq_X_, range(len(unq_X)))
+        )
 
         out = np.empty(
-            (len(X), vocabulary_count_matrix.shape[0]), dtype=self.dtype)
+            (len(X), vocabulary_count_matrix.shape[0]),
+            dtype=self.dtype,
+        )
 
         for x, out_row in zip(X, out):
             out_row[:] = se_dict[x]
 
         return np.nan_to_num(out, copy=False)
 
-    def fit_transform(self, X, y=None, **fit_params):
-            """
-            Fit SimilarityEncoder to data, then transform it.
-            Fits transformer to `X` and `y` with optional parameters
-            `fit_params` and returns a transformed version of `X`.
-
-            Parameters
-            ----------
-            X : array-like of shape (n_samples, n_features)
-                Input samples.
-            y :  array-like of shape (n_samples,) or (n_samples, n_outputs), \
-                    default=None
-                Target values (None for unsupervised transformations).
-            **fit_params : dict
-                Additional fit parameters.
-
-            Returns
-            -------
-            X_new : ndarray array of shape (n_samples, n_features_new)
-                Transformed array.
-            """
-            if y is None:
-                # fit method of arity 1 (unsupervised transformation)
-                return self.fit(X, **fit_params).transform(X)
-            else:
-                # fit method of arity 2 (supervised transformation)
-                return self.fit(X, y, **fit_params).transform(X)
+    def fit_transform(self, X, y=None, **fit_params) -> np.array:
+        """
+        Fit SimilarityEncoder to data, then transform it.
+        Fits transformer to `X` and `y` with optional parameters
+        `fit_params` and returns a transformed version of `X`.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input samples.
+        y : array-like of shape (n_samples,) or (n_samples, n_outputs), \
+                default=None
+            Target values (None for unsupervised transformations).
+        **fit_params
+            Additional fit parameters.
+
+        Returns
+        -------
+        array of shape (n_samples, n_features_new)
+            Transformed array.
+        """
+        if y is None:
+            # fit method of arity 1 (unsupervised transformation)
+            return self.fit(X, **fit_params).transform(X)
+        else:
+            # fit method of arity 2 (supervised transformation)
+            return self.fit(X, y, **fit_params).transform(X)
diff --git a/dirty_cat/string_distances.py b/dirty_cat/string_distances.py
index 80c4189..aefc68b 100644
--- a/dirty_cat/string_distances.py
+++ b/dirty_cat/string_distances.py
@@ -3,16 +3,17 @@ Some string distances
 """
 
 import re
-import numpy as np
 
+from typing import Tuple
 from collections import Counter
 
 
 # TODO vectorize these functions (accept arrays)
 
 
-def get_ngram_count(X, ngram_range):
-    """Compute the number of ngrams in a string.
+def get_ngram_count(string: str, ngram_range: Tuple[int, int]) -> int:
+    """
+    Compute the number of ngrams in a string.
 
     Here is where the formula comes from:
 
@@ -27,17 +28,18 @@ def get_ngram_count(X, ngram_range):
     ngram_count = 0
 
     for i in range(min_n, max_n + 1):
-        ngram_count += len(X) - i + 1
+        ngram_count += len(string) - i + 1
 
     return ngram_count
 
 
-def preprocess(x):
-    """Combine preprocessing done by CountVectorizer and the SimilarityEncoder.
+def preprocess(x: str) -> str:
+    """
+    Combine preprocessing done by CountVectorizer and the SimilarityEncoder.
 
     Different methods exist to compute the number of ngrams in a string:
 
-    - Simply sum the values of a count vector, which is the ouput of a
+    - Simply sum the values of a count vector, which is the output of a
       CountVectorizer with analyzer="char", and a specific ngram_range
     - Compute the number of ngrams using a formula (see ``get_ngram_count``)
 
@@ -48,22 +50,39 @@ def preprocess(x):
     prior to the CountVectorizer.
     """
 
-    # preprocessing step done in ngram_similarity
-    x = ' %s ' % x
+    # Preprocessing step done in ngram_similarity
+    x = f' {x} '
 
-    # preprocessing step done in the CountVectorizer
+    # Preprocessing step done in the CountVectorizer
     _white_spaces = re.compile(r"\s\s+")
 
     return _white_spaces.sub(' ', x)
 
 
-def get_unique_ngrams(string, n):
-    """ Return the set of different tri-grams in a string
+def get_unique_ngrams(string: str, ngram_range: Tuple[int, int]):
+    """
+    Return the set of unique n-grams of a string.
+
+    Parameters
+    ----------
+    string : str
+        The string to split in n-grams.
+    ngram_range : tuple (min_n, max_n)
+    The lower and upper boundary of the range of n-values for different
+    n-grams to be extracted. All values of n such that min_n <= n <= max_n.
+
+    Returns
+    -------
+    set
+        The set of unique n-grams of the string.
     """
     spaces = ' '  # * (n // 2 + n % 2)
     string = spaces + " ".join(string.lower().split()) + spaces
-    string_list = [string[i:] for i in range(n)]
-    return set(zip(*string_list))
+    ngram_set = set()
+    for n in range(ngram_range[0], ngram_range[1] + 1):
+        string_list = [string[i:] for i in range(n)]
+        ngram_set |= set(zip(*string_list))
+    return ngram_set
 
 
 def get_ngrams(string, n):
diff --git a/dirty_cat/super_vectorizer.py b/dirty_cat/super_vectorizer.py
index e56b0cb..35bd4e0 100644
--- a/dirty_cat/super_vectorizer.py
+++ b/dirty_cat/super_vectorizer.py
@@ -1,20 +1,14 @@
 """
-
-This class implements the SuperVectorizer, which is a preprocessor used to
-automatically apply encoders to different types of data, without the need to
-manually categorize them beforehand, or construct complex Pipelines.
-
+Implements the SuperVectorizer: a preprocessor to automatically
+apply encoders to different types of data, without the need to manually
+categorize them beforehand, or construct complex Pipelines.
 """
 
-# Author: Lilian Boulard <lilian@boulard.fr> | https://github.com/LilianBoulard
-
-import sklearn
-
 import numpy as np
 import pandas as pd
 
 from warnings import warn
-from typing import Union, Optional, List
+from typing import Union, Optional, List, Dict, Tuple, Literal
 
 from sklearn.base import BaseEstimator, clone
 from sklearn.compose import ColumnTransformer
@@ -56,7 +50,7 @@ class SuperVectorizer(ColumnTransformer):
     Parameters
     ----------
 
-    cardinality_threshold: int, default=40
+    cardinality_threshold : int, default=40
         Two lists of features will be created depending on this value: strictly
         under this value, the low cardinality categorical values, and above or
         equal, the high cardinality categorical values.
@@ -64,7 +58,7 @@ class SuperVectorizer(ColumnTransformer):
         the parameters `low_card_cat_transformer` and
         `high_card_cat_transformer` respectively.
 
-    low_card_cat_transformer: Transformer or str or None, default=None
+    low_card_cat_transformer : typing.Optional[typing.Union[sklearn.base.BaseEstimator, str]], default=None
         Transformer used on categorical/string features with low cardinality
         (threshold is defined by `cardinality_threshold`).
         Default value None is converted to `OneHotEncoder()`.
@@ -73,7 +67,7 @@ class SuperVectorizer(ColumnTransformer):
         None to apply `remainder`, 'drop' for dropping the columns,
         or 'passthrough' to return the unencoded columns.
 
-    high_card_cat_transformer: Transformer or str or None, default=None
+    high_card_cat_transformer : typing.Optional[typing.Union[sklearn.base.BaseEstimator, str]], default=None
         Transformer used on categorical/string features with high cardinality
         (threshold is defined by `cardinality_threshold`).
         Default value None is converted to `GapEncoder(n_components=30)`.
@@ -82,14 +76,14 @@ class SuperVectorizer(ColumnTransformer):
         None to apply `remainder`, 'drop' for dropping the columns,
         or 'passthrough' to return the unencoded columns.
 
-    numerical_transformer: Transformer or str or None, default=None
+    numerical_transformer : typing.Optional[typing.Union[sklearn.base.BaseEstimator, str]], default=None
         Transformer used on numerical features.
         Can either be a transformer object instance (e.g. `StandardScaler()`),
         a `Pipeline` containing the preprocessing steps,
         None to apply `remainder`, 'drop' for dropping the columns,
         or 'passthrough' to return the unencoded columns.
 
-    datetime_transformer: Transformer or str or None, default=None
+    datetime_transformer : typing.Optional[typing.Union[sklearn.base.BaseEstimator, str]], default=None
         Transformer used on datetime features.
         Default value None is converted to `DatetimeEncoder()`.
         Can either be a transformer object instance (e.g. `DatetimeEncoder()`),
@@ -97,11 +91,11 @@ class SuperVectorizer(ColumnTransformer):
         None to apply `remainder`, 'drop' for dropping the columns,
         or 'passthrough' to return the unencoded columns.
 
-    auto_cast: bool, default=True
+    auto_cast : bool, default=True
         If set to `True`, will try to convert each column to the best possible
         data type (dtype).
 
-    impute_missing: str, default='auto'
+    impute_missing : str, default='auto'
         When to impute missing values in string columns.
         'auto' will impute missing values if it's considered appropriate
         (we are using an encoder that does not support missing values and/or
@@ -111,7 +105,7 @@ class SuperVectorizer(ColumnTransformer):
         When imputed, missing values are replaced by the string 'missing'.
         See also attribute `imputed_columns_`.
 
-    remainder : {'drop', 'passthrough'} or estimator, default='drop'
+    remainder : typing.Union[typing.Literal["drop", "passthrough"], sklearn.base.BaseEstimator], default='drop'
         By default, only the specified columns in `transformers` are
         transformed and combined in the output, and the non-specified
         columns are dropped. (default of ``'drop'``).
@@ -129,8 +123,8 @@ class SuperVectorizer(ColumnTransformer):
         If the output of the different transformers contains sparse matrices,
         these will be stacked as a sparse matrix if the overall density is
         lower than this value. Use sparse_threshold=0 to always return dense.
-        When the transformed output consists of all dense data, the stacked result
-        will be dense, and this keyword will be ignored.
+        When the transformed output consists of all dense data, the stacked
+        result will be dense, and this keyword will be ignored.
 
     n_jobs : int, default=None
         Number of jobs to run in parallel.
@@ -149,26 +143,31 @@ class SuperVectorizer(ColumnTransformer):
     Attributes
     ----------
 
-    transformers_: List[Tuple[str, Union[str, BaseEstimator], Union[str, int]]]
+    transformers_: typing.List[typing.Tuple[str, typing.Union[str, sklearn.base.BaseEstimator], typing.List[str]]]
         The final distribution of columns.
         List of three-tuple containing
         (1) the name of the category
         (2) the encoder/transformer instance which will be applied
         or "passthrough" or "drop"
-        (3) the list of column names or index
+        (3) the list of column names
 
-    columns_: List[Union[str, int]]
+    columns_: typing.List[typing.Union[str, int]]
         The column names of fitted array.
 
-    types_: Dict[str, type]
+    types_: typing.Dict[str, type]
         A mapping of inferred types per column.
         Key is the column name, value is the inferred dtype.
 
-    imputed_columns_: List[str]
+    imputed_columns_: typing.List[str]
         The list of columns in which we imputed the missing values.
 
     """
 
+    transformers_: List[Tuple[str, Union[str, BaseEstimator], List[str]]]
+    columns_: List[str]
+    types_: Dict[str, type]
+    imputed_columns_: List[str]
+
     # Override required parameters
     _required_parameters = []
     OptionalEstimator = Optional[Union[BaseEstimator, str]]
@@ -181,8 +180,8 @@ class SuperVectorizer(ColumnTransformer):
                  datetime_transformer: Optional[Union[BaseEstimator, str]] = None,
                  auto_cast: bool = True,
                  impute_missing: str = 'auto',
-                 # Following parameters are inherited from ColumnTransformer
-                 remainder: str = 'passthrough',
+                 # The next parameters are inherited from ColumnTransformer
+                 remainder: Union[Literal["drop", "passthrough"], BaseEstimator] = 'passthrough',
                  sparse_threshold: float = 0.3,
                  n_jobs: int = None,
                  transformer_weights=None,
@@ -224,7 +223,7 @@ class SuperVectorizer(ColumnTransformer):
         else:
             self.datetime_transformer_ = DatetimeEncoder()
 
-        #TODO check that the provided transformers are valid
+        # TODO: check that the provided transformers are valid
 
     @staticmethod
     def _auto_cast(X: pd.DataFrame) -> pd.DataFrame:
@@ -240,8 +239,8 @@ class SuperVectorizer(ColumnTransformer):
         Returns
         -------
         pd.DataFrame
-            The same pandas DataFrame, with its columns casted to the best possible
-            data type.
+            The same pandas DataFrame, with its columns cast to the best
+            possible data type.
         """
         from pandas.core.dtypes.base import ExtensionDtype
 
@@ -253,25 +252,29 @@ class SuperVectorizer(ColumnTransformer):
             # with scikit-learn (as of version 0.24).
             if contains_missing:
                 # Some numerical dtypes like Int64 or Float64 only support
-                # pd.NA so they must be converted to np.float64 before.
+                # pd.NA, so they must be converted to np.float64 before.
                 if pd.api.types.is_numeric_dtype(X[col]):
-                    X[col] = X[col].astype(np.float64)
+                    X.loc[:, col] = X[col].astype(np.float64)
                 X[col].fillna(value=np.nan, inplace=True)
-        STR_NA_VALUES = ['null', '', '1.#QNAN', '#NA', 'nan', '#N/A N/A', '-1.#QNAN', '<NA>', '-1.#IND', '-nan', 'n/a',
-                         '-NaN', '1.#IND', 'NULL', 'NA', 'N/A', '#N/A', 'NaN']  # taken from pandas.io.parsers (version 1.1.4)
-        X = X.replace(STR_NA_VALUES + [None, "?", "..."],
-                      np.nan)
-        X = X.replace(r'^\s+$', np.nan, regex=True) # replace whitespace only
 
-        # Convert to best possible data type
+        # taken from pandas.io.parsers (version 1.1.4)
+        STR_NA_VALUES = ['null', '', '1.#QNAN', '#NA', 'nan', '#N/A N/A',
+                         '-1.#QNAN', '<NA>', '-1.#IND', '-nan', 'n/a',
+                         '-NaN', '1.#IND', 'NULL', 'NA', 'N/A', '#N/A', 'NaN']
+        X = X.replace(STR_NA_VALUES + [None, "?", "..."], np.nan)
+        X = X.replace(r'^\s+$', np.nan, regex=True)  # replace whitespaces
+
+        # Convert to the best possible data type
         for col in X.columns:
-            if not pd.api.types.is_datetime64_any_dtype(X[col]): # we don't want to cast datetime64
+            if not pd.api.types.is_datetime64_any_dtype(X[col]):
+                # we don't want to cast datetime64
                 try:
-                    X[col] = pd.to_numeric(X[col], errors='raise')
+                    X.loc[:, col] = pd.to_numeric(X[col], errors='raise')
                 except:
-                    # Only try to convert to datetime if the variable isn't numeric.
+                    # Only try to convert to datetime
+                    # if the variable isn't numeric.
                     try:
-                        X[col] = pd.to_datetime(X[col], errors='raise',
+                        X.loc[:, col] = pd.to_datetime(X[col], errors='raise',
                                                 infer_datetime_format=True)
                     except:
                         pass
@@ -279,13 +282,16 @@ class SuperVectorizer(ColumnTransformer):
             # for earlier versions of sklearn
             if issubclass(X[col].dtype.__class__, ExtensionDtype):
                 try:
-                    X[col] = X[col].astype(X[col].dtype.type, errors='ignore')
+                    X.loc[:, col] = X[col].astype(X[col].dtype.type,
+                                                  errors='ignore')
                 except (TypeError, ValueError):
                     pass
         return X
 
     def transform(self, X) -> np.ndarray:
-        """Transform X by applying transformers on each column, then concatenate results.
+        """
+        Transform X by applying fitted transformers on each column,
+        and concatenate the results.
 
         Parameters
         ----------
@@ -294,13 +300,11 @@ class SuperVectorizer(ColumnTransformer):
 
         Returns
         -------
-        X_t : {array-like, sparse matrix} of \
-                shape (n_samples, sum_n_components)
+        {array-like, sparse matrix} of shape (n_samples, sum_n_components)
             hstack of results of transformers. sum_n_components is the
             sum of n_components (output dimension) over transformers. If
             any result is a sparse matrix, everything will be converted to
             sparse matrices.
-
         """
         if not isinstance(X, pd.DataFrame):
             X = pd.DataFrame(X)
@@ -313,7 +317,7 @@ class SuperVectorizer(ColumnTransformer):
             X = self._auto_cast(X)
             self.types_ = {c: t for c, t in zip(X.columns, X.dtypes)}
         if X.shape[1] != len(self.columns_):
-            raise ValueError("Passed array does not match column count of "
+            raise ValueError(f"Passed array does not match column count of "
                              f"array seen at fit time. Got {X.shape[1]} "
                              f"columns, expected {len(self.columns_)}")
 
@@ -323,26 +327,25 @@ class SuperVectorizer(ColumnTransformer):
             X.columns = self.columns_
 
         for col in self.imputed_columns_:
-            X[col] = _replace_missing_in_col(X[col])
+            X.loc[:, col] = _replace_missing_in_col(X[col])
 
         return super().transform(X)
 
     def fit_transform(self, X, y=None):
         """
-        Fit all transformers, transform the data, and concatenate results.
+        Fit all transformers, transform the data, and concatenate the results.
 
         Parameters
         ----------
-        X: {array-like, dataframe} of shape (n_samples, n_features)
+        X : {array-like, dataframe} of shape (n_samples, n_features)
             Input data, of which specified subsets are used to fit the
             transformers.
-
-        y: array-like of shape (n_samples,), default=None
+        y : array-like of shape (n_samples,), default=None
             Targets for supervised learning.
 
         Returns
         -------
-        X_t: {array-like, sparse matrix} of shape (n_samples, sum_n_components)
+        {array-like, sparse matrix} of shape (n_samples, sum_n_components)
             hstack of results of transformers. sum_n_components is the
             sum of n_components (output dimension) over transformers. If
             any result is a sparse matrix, everything will be converted to
@@ -363,21 +366,26 @@ class SuperVectorizer(ColumnTransformer):
             # Create a copy to avoid altering the original data.
             X = X.copy()
 
-        self.columns_ = X.columns
+        self.columns_ = X.columns.to_list()
         # If auto_cast is True, we'll find and apply the best possible type
         # to each column.
-        # We'll keep the results so we can apply the types in transform.
+        # We'll keep the results in order to apply the types in `transform`.
         if self.auto_cast:
             X = self._auto_cast(X)
             self.types_ = {c: t for c, t in zip(X.columns, X.dtypes)}
 
         # Select columns by dtype
-        numeric_columns = X.select_dtypes(include=['int', 'float',
-                                                   np.float64, np.float32, np.float16,
-                                                   np.int64, np.int32, np.int16,
-                                                   np.uint64, np.uint32, np.uint16]).columns.to_list()
-        categorical_columns = X.select_dtypes(include=['string', 'object', 'category']).columns.to_list()
-        datetime_columns = X.select_dtypes(include=['datetime', "datetimetz"]).columns.to_list()
+        numeric_columns = X.select_dtypes(
+            include=[
+                'int', 'float',
+                np.float64, np.float32, np.float16,
+                np.int64, np.int32, np.int16,
+                np.uint64, np.uint32, np.uint16
+            ]).columns.to_list()
+        categorical_columns = X.select_dtypes(
+            include=['string', 'object', 'category']).columns.to_list()
+        datetime_columns = X.select_dtypes(
+            include=['datetime', "datetimetz"]).columns.to_list()
 
         # Divide categorical columns by cardinality
         low_card_cat_columns = [
@@ -407,7 +415,7 @@ class SuperVectorizer(ColumnTransformer):
                 self.transformers.append(trans)
 
         if len(self.transformers) == 0:
-            raise RuntimeError('No transformers could be generated !')
+            raise RuntimeError('No transformers could be generated! ')
 
         self.imputed_columns_ = []
         if _has_missing_values(X):
@@ -416,7 +424,7 @@ class SuperVectorizer(ColumnTransformer):
                 for col in X.columns:
                     # Do not impute numeric columns
                     if not pd.api.types.is_numeric_dtype(X[col]):
-                        X[col] = _replace_missing_in_col(X[col])
+                        X.loc[:, col] = _replace_missing_in_col(X[col])
                         self.imputed_columns_.append(col)
 
             elif self.impute_missing == 'skip':
@@ -424,10 +432,6 @@ class SuperVectorizer(ColumnTransformer):
 
             elif self.impute_missing == 'auto':
                 for name, trans, cols in all_transformers:
-                    # At each iteration, we'll manipulate a boolean,
-                    # and depending on its value at the end of the loop,
-                    # we will or will not replace the missing values in
-                    # the columns.
                     impute: bool = False
 
                     if isinstance(trans, OneHotEncoder) \
@@ -438,35 +442,41 @@ class SuperVectorizer(ColumnTransformer):
                         for col in cols:
                             # Do not impute numeric columns
                             if not pd.api.types.is_numeric_dtype(X[col]):
-                                X[col] = _replace_missing_in_col(X[col])
+                                X.loc[:, col] = _replace_missing_in_col(X[col])
                                 self.imputed_columns_.append(col)
 
             else:
                 raise ValueError(
                     "Invalid value for `impute_missing`, expected any of "
-                    f"{'auto', 'force', 'skip'}, got {self.impute_missing!r}."
+                    f"{{'auto', 'force', 'skip'}}, "
+                    f"got {self.impute_missing!r}. "
                 )
 
         # If there was missing values imputation, we cast the DataFrame again,
-        # as pandas give different types depending whether a column has
+        # as pandas gives different types depending on whether a column has
         # missing values or not.
         if self.imputed_columns_:
             X = self._auto_cast(X)
 
         if self.verbose:
-            print(f'[SuperVectorizer] Assigned transformers: {self.transformers}')
+            print(f'[SuperVectorizer] Assigned transformers: '
+                  f'{self.transformers}')
 
         check_input(X)
 
         res = super().fit_transform(X, y)
 
-        # for the "remainder" columns, the ColumnTransformer transformers_ attribute
-        # contains the index instead of the column name, so we convert it to the column name
+        # For the "remainder" columns, the `ColumnTransformer` `transformers_`
+        # attribute contains the index instead of the column name,
+        # so we convert the values to the appropriate column names
         # if there is less than 20 columns in the remainder.
-        for i, tup in enumerate(self.transformers_):
-            name, enc, cols = tup  # Unpack
+        for i, (name, enc, cols) in enumerate(self.transformers_):
             if name == "remainder" and len(cols) < 20:
-                self.transformers_[i] = (name, enc, [self.columns_[j] for j in cols])
+                # In this case, "cols" is a list of ints (the indices)
+                cols: List[int]
+                self.transformers_[i] = (
+                    name, enc, [self.columns_[j] for j in cols]
+                )
 
         return res
 
@@ -474,8 +484,12 @@ class SuperVectorizer(ColumnTransformer):
         """
         Returns clean feature names with format
         "<column_name>_<value>" if encoded by OneHotEncoder or alike,
-        e.g. "job_title_Police officer",
-        or "<column_name>" if not encoded.
+        e.g. "job_title_Police officer", or "<column_name>" otherwise.
+
+        Returns
+        -------
+        typing.List[str]
+            Feature names.
         """
         if Version(sklearn_version) < Version('0.23'):
             try:
@@ -488,7 +502,7 @@ class SuperVectorizer(ColumnTransformer):
                     'Prior to sklearn 0.23, get_feature_names with '
                     '"passthrough" is unsupported. To use the method, '
                     'either make sure there is no "passthrough" in the '
-                    'transformers, or update your copy of scikit-learn.'
+                    'transformers, or update your copy of scikit-learn. '
                 )
         else:
             if Version(sklearn_version) < Version('1.0'):
@@ -516,12 +530,12 @@ class SuperVectorizer(ColumnTransformer):
                 all_trans_feature_names.extend(trans_feature_names)
 
         if len(ct_feature_names) != len(all_trans_feature_names):
-            warn('Could not extract clean feature names ; returning defaults.')
+            warn('Could not extract clean feature names; returning defaults. ')
             return ct_feature_names
 
         return all_trans_feature_names
-    
-    def get_feature_names(self) -> List[str]:
+
+    def get_feature_names(self, input_features=None) -> List[str]:
         """
         Ensures compatibility with sklearn < 1.0.
         Use `get_feature_names_out` instead.
@@ -530,8 +544,8 @@ class SuperVectorizer(ColumnTransformer):
             warn(
                 "Following the changes in scikit-learn 1.0, "
                 "get_feature_names is deprecated. "
-                "Use get_feature_names_out instead.",
+                "Use get_feature_names_out instead. ",
                 DeprecationWarning,
                 stacklevel=2,
             )
-        return self.get_feature_names_out()
+        return self.get_feature_names_out(input_features)
diff --git a/dirty_cat/target_encoder.py b/dirty_cat/target_encoder.py
index fc90950..7b0622b 100644
--- a/dirty_cat/target_encoder.py
+++ b/dirty_cat/target_encoder.py
@@ -1,6 +1,8 @@
 import collections
 import numpy as np
 
+from typing import List, Dict, Literal, Union
+
 from sklearn.preprocessing import LabelEncoder
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils import check_array
@@ -10,8 +12,7 @@ from dirty_cat.utils import check_input
 
 
 def lambda_(x, n):
-    out = x / (x + n)
-    return out
+    return x / (x + n)
 
 
 class TargetEncoder(BaseEstimator, TransformerMixin):
@@ -21,35 +22,29 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
     target variable y. The method considers that categorical
     variables can present rare categories. It represents each category by the
     probability of y conditional on this category.
-    In addition it takes an empirical Bayes approach to shrink the estimate.
+    In addition, it takes an empirical Bayes approach to shrink the estimate.
 
 
     Parameters
     ----------
-    categories : 'auto' or a list of lists/arrays of values.
+    categories : typing.Union[typing.Literal["auto"], typing.List[typing.List[typing.Union[str, int]]]
         Categories (unique values) per feature:
-
         - 'auto' : Determine categories automatically from the training data.
         - list : ``categories[i]`` holds the categories expected in the i-th
           column. The passed categories must be sorted and should not mix
           strings and numeric values.
-
         The categories used can be found in the ``categories_`` attribute.
-
-    clf_type : string {'regression', 'binary-clf', 'multiclass-clf'}
+    clf_type : typing.Literal["regression", "binary-clf", "multiclass-clf"]
         The type of classification/regression problem.
-
-    dtype : number type, default np.float64
+    dtype : type, default=np.float64
         Desired dtype of output.
-
-    handle_unknown : 'error' (default) or 'ignore'
+    handle_unknown : typing.Literal["error", "ignore"], default="error"
         Whether to raise an error or ignore if a unknown categorical feature is
         present during transform (default is to raise). When this parameter
         is set to 'ignore' and an unknown category is encountered during
         transform, the resulting one-hot encoded columns for this feature
         will be all zeros.
-
-    handle_missing : 'error' or '' (default)
+    handle_missing : typing.Literal["error", ""], default=""
         Whether to raise an error or impute with blank string '' if missing
         values (NaN) are present during fit (default is to impute).
         When this parameter is set to '', and a missing value is encountered
@@ -58,7 +53,9 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
 
     Attributes
     ----------
-    categories_ : list of arrays
+    n_features_in_: int
+        Number of features in the data seen during fit.
+    categories_ : typing.List[np.array]
         The categories of each feature determined during fitting
         (in order corresponding with output of ``transform``).
 
@@ -68,25 +65,34 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
     high-cardinality categorical attributes in classification and prediction
     problems.
     """
+
+    n_features_in_: int
+    _label_encoders_: List[LabelEncoder]
+    categories_: List[np.array]
+    n_: int
+
     def __init__(self,
-                 categories='auto',
-                 clf_type='binary-clf',
-                 dtype=np.float64, handle_unknown='error',
-                 handle_missing=''):
+                 categories: Union[Literal["auto"], List[Union[List[str], np.array]]] = 'auto',
+                 clf_type: Literal["regression", "binary-clf", "multiclass-clf"] = 'binary-clf',
+                 dtype: type = np.float64,
+                 handle_unknown: Literal["error", "ignore"] = 'error',
+                 handle_missing: Literal["error", ""] = ''):
         self.categories = categories
         self.dtype = dtype
         self.clf_type = clf_type
         self.handle_unknown = handle_unknown
         self.handle_missing = handle_missing
 
-    def _more_tags(self):
+    def _more_tags(self) -> Dict[str, List[str]]:
         """
         Used internally by sklearn to ease the estimator checks.
         """
         return {"X_types": ["categorical"]}
 
-    def fit(self, X, y):
-        """Fit the TargetEncoder to X.
+    def fit(self, X, y) -> "TargetEncoder":
+        """
+        Fit the TargetEncoder to X.
+
         Parameters
         ----------
         X : array-like, shape [n_samples, n_features]
@@ -96,20 +102,23 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
 
         Returns
         -------
-        self
+        TargetEncoder
+            Fitted TargetEncoder instance.
         """
         X = check_input(X)
         self.n_features_in_ = X.shape[1]
         if self.handle_missing not in ['error', '']:
-            template = ("handle_missing should be either 'error' or "
-                        "'', got %s")
-            raise ValueError(template % self.handle_missing)
+            raise ValueError(
+                f"Got handle_missing={self.handle_missing!r}, but expected "
+                f"any of {{'error', ''}}. "
+            )
         if hasattr(X, 'iloc') and X.isna().values.any():
             if self.handle_missing == 'error':
-                msg = ("Found missing values in input data; set "
-                       "handle_missing='' to encode with missing values")
-                raise ValueError(msg)
-            if self.handle_missing != 'error':
+                    raise ValueError(
+                        "Found missing values in input data; set "
+                        "handle_missing='' to encode with missing values. "
+                    )
+            else:
                 X = X.fillna(self.handle_missing)
         elif not hasattr(X, 'dtype') and isinstance(X, list):
             X = np.asarray(X, dtype=object)
@@ -118,28 +127,32 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
             mask = _object_dtype_isnan(X)
             if X.dtype.kind == 'O' and mask.any():
                 if self.handle_missing == 'error':
-                    msg = ("Found missing values in input data; set "
-                           "handle_missing='' to encode with missing values")
-                    raise ValueError(msg)
-                if self.handle_missing != 'error':
+                    raise ValueError(
+                        "Found missing values in input data; set "
+                        "handle_missing='' to encode with missing values. "
+                    )
+                else:
                     X[mask] = self.handle_missing
 
         if self.handle_unknown not in ['error', 'ignore']:
-            template = ("handle_unknown should be either 'error' or "
-                        "'ignore', got %s")
-            raise ValueError(template % self.handle_unknown)
+            raise ValueError(
+                f'Got handle_unknown={self.handle_unknown!r}, but expected'
+                f'any of {{"error", "ignore"}}. '
+            )
 
         if self.categories != 'auto':
             for cats in self.categories:
                 if not np.all(np.sort(cats) == np.array(cats)):
-                    raise ValueError("Unsorted categories are not yet "
-                                     "supported")
+                    raise ValueError(
+                        "Unsorted categories are not yet supported. "
+                    )
 
         X_temp = check_array(X, dtype=None)
         if not hasattr(X, 'dtype') and np.issubdtype(X_temp.dtype, np.str_):
             X = check_array(X, dtype=np.object)
         else:
             X = X_temp
+        X: np.array
 
         n_samples, n_features = X.shape
 
@@ -155,9 +168,10 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
                     valid_mask = np.in1d(Xj, self.categories[j])
                     if not np.all(valid_mask):
                         diff = np.unique(Xj[~valid_mask])
-                        msg = ("Found unknown categories {0} in column {1}"
-                               " during fit".format(diff, j))
-                        raise ValueError(msg)
+                        raise ValueError(
+                            f"Found unknown categories {diff} in column {j}"
+                            f" during fit"
+                        )
                 le.classes_ = np.array(self.categories[j])
 
         self.categories_ = [le.classes_ for le in self._label_encoders_]
@@ -182,8 +196,9 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
         self.k_ = {j: len(self.counter_[j]) for j in self.counter_}
         return self
 
-    def transform(self, X):
-        """Transform X using specified encoding scheme.
+    def transform(self, X) -> np.array:
+        """
+        Transform X using the specified encoding scheme.
 
         Parameters
         ----------
@@ -192,21 +207,23 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
 
         Returns
         -------
-        X_new : 2-d array
+        2-d np.array
             Transformed input.
         """
         check_is_fitted(self, attributes=["n_features_in_"])
         X = check_input(X)
         if X.shape[1] != self.n_features_in_:
             raise ValueError(
-                f"Number of features in the input data ({X.shape[1]}) does not match the number of features "
+                f"The number of features in the input data ({X.shape[1]}) "
+                f"does not match the number of features "
                 f"seen during fit ({self.n_features_in_})."
             )
         if hasattr(X, 'iloc') and X.isna().values.any():
             if self.handle_missing == 'error':
-                msg = ("Found missing values in input data; set "
-                       "handle_missing='' to encode with missing values")
-                raise ValueError(msg)
+                    raise ValueError(
+                        "Found missing values in input data; set "
+                        "handle_missing='' to encode with missing values. "
+                    )
             if self.handle_missing != 'error':
                 X = X.fillna(self.handle_missing)
         elif not hasattr(X, 'dtype') and isinstance(X, list):
@@ -216,9 +233,10 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
             mask = _object_dtype_isnan(X)
             if X.dtype.kind == 'O' and mask.any():
                 if self.handle_missing == 'error':
-                    msg = ("Found missing values in input data; set "
-                           "handle_missing='' to encode with missing values")
-                    raise ValueError(msg)
+                    raise ValueError(
+                        "Found missing values in input data; set "
+                        "handle_missing='' to encode with missing values. "
+                    )
                 if self.handle_missing != 'error':
                     X[mask] = self.handle_missing
 
@@ -239,9 +257,10 @@ class TargetEncoder(BaseEstimator, TransformerMixin):
             if not np.all(valid_mask):
                 if self.handle_unknown == 'error':
                     diff = np.unique(X[~valid_mask, i])
-                    msg = ("Found unknown categories {0} in column {1}"
-                           " during transform".format(diff, i))
-                    raise ValueError(msg)
+                    raise ValueError(
+                        f"Found unknown categories {diff} in column {i} "
+                        f"during transform. "
+                    )
                 else:
                     # Set the problematic rows to an acceptable value and
                     # continue `The rows are marked `X_mask` and will be
diff --git a/dirty_cat/utils.py b/dirty_cat/utils.py
index 1a22d85..ef4e80a 100644
--- a/dirty_cat/utils.py
+++ b/dirty_cat/utils.py
@@ -1,21 +1,19 @@
 import collections
-
 import numpy as np
-import pandas as pd
-from sklearn.utils import check_X_y, check_array
 
-from typing import Tuple, Union
+from sklearn.utils import check_array
+from typing import Tuple, Union, Any, Hashable
 
 
 class LRUDict:
     """ dict with limited capacity
 
-    Using LRU eviction, this avoid to memorizz a full dataset"""
-    def __init__(self, capacity):
+    Using LRU eviction avoids memorizing a full dataset"""
+    def __init__(self, capacity: int):
         self.capacity = capacity
         self.cache = collections.OrderedDict()
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Hashable):
         try:
             value = self.cache.pop(key)
             self.cache[key] = value
@@ -23,7 +21,7 @@ class LRUDict:
         except KeyError:
             return -1
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: Hashable, value: Any):
         try:
             self.cache.pop(key)
         except KeyError:
@@ -31,36 +29,41 @@ class LRUDict:
                 self.cache.popitem(last=False)
         self.cache[key] = value
 
-    def __contains__(self, key):
+    def __contains__(self, key: Hashable):
         return key in self.cache
 
 
-def check_input(X):
+def check_input(X) -> np.array:
     """
     Check input with sklearn standards.
     Also converts X to a numpy array if not already.
     """
-    #TODO check for weird type of input to pass scikit learn tests
-    # whithout messing with the original type too much
-
-    X_ = check_array(X,
-                dtype= None,
-                ensure_2d=True,
-                force_all_finite=False)
-    #If the array contains both NaNs and strings, convert to object type
-    if X_.dtype.kind in {'U', 'S'}: # contains strings
-        if np.any(X_ == "nan"): # missing value converted to string
-            return check_array(np.array(X, dtype=object),
-                dtype= None,
+    # TODO check for weird type of input to pass scikit learn tests
+    #  without messing with the original type too much
+
+    X_ = check_array(
+        X,
+        dtype=None,
+        ensure_2d=True,
+        force_all_finite=False,
+    )
+    # If the array contains both NaNs and strings, convert to object type
+    if X_.dtype.kind in {'U', 'S'}:  # contains strings
+        if np.any(X_ == "nan"):  # missing value converted to string
+            return check_array(
+                np.array(X, dtype=object),
+                dtype=None,
                 ensure_2d=True,
-                force_all_finite=False)
+                force_all_finite=False,
+            )
 
     return X_
 
 
 class Version:
     """
-    Replacement for `distutil.version.LooseVersion` and `packaging.version.Version`.
+    Replacement for `distutil.version.LooseVersion` and
+    `packaging.version.Version`.
     Implemented to avoid `DeprecationWarning`s raised by the former,
     and avoid adding a dependency for the latter.
 
@@ -81,7 +84,7 @@ class Version:
     >>> # You can also pass the separator for specific cases
     >>> Version('1-5', separator='-') == Version('1-6-5', separator='-')
     >>> Version('1-5', separator='-') == '1-6-5'
-    >>> Version('1-5', separator='-') == '1.6.5'  # Won't work !
+    >>> Version('1-5', separator='-') == '1.6.5'  # Won't work!
     """
 
     def __init__(self, value: str, separator: str = '.'):
diff --git a/doc/_templates/class.rst b/doc/_templates/class.rst
index 4a1feb7..7ddd0bc 100644
--- a/doc/_templates/class.rst
+++ b/doc/_templates/class.rst
@@ -9,7 +9,6 @@
 
 .. autoclass:: {{ objname }}
    :inherited-members:
-   :undoc-members:
 
    {% block methods %}
 
diff --git a/doc/conf.py b/doc/conf.py
index 36042e6..1afcf7d 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -28,17 +28,18 @@ import os
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = ['sphinx.ext.autodoc',
-              'sphinx.ext.autosummary',
-              'sphinx.ext.doctest',
-              'sphinx.ext.intersphinx',
-              'sphinx.ext.mathjax',
-              'sphinx.ext.viewcode',
-              'sphinx.ext.githubpages',
-              'sphinx.ext.napoleon',
-              'sphinx.ext.autodoc.typehints',
-              'sphinx_gallery.gen_gallery',
-              ]
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.autosummary',
+    'sphinx.ext.doctest',
+    'sphinx.ext.intersphinx',
+    'sphinx.ext.mathjax',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.githubpages',
+     'sphinx.ext.napoleon',
+    'sphinx.ext.autodoc.typehints',
+    'sphinx_gallery.gen_gallery',
+]
 
 try:
     import sphinxext.opengraph
@@ -68,8 +69,7 @@ author = u'dirty_cat developers'
 # built documents.
 #
 # The short X.Y version.
-version_file = os.path.join(
-    '..', 'dirty_cat', 'VERSION.txt')
+version_file = os.path.join('..', 'dirty_cat', 'VERSION.txt')
 with open(version_file) as fh:
     version = fh.read().strip()
 # The full version, including alpha/beta/rc tags.
@@ -100,12 +100,9 @@ autosummary_generate = True
 # Add any paths that contain templates here, relative to this directory.
 #templates_path = ['_templates']
 
-
 autodoc_default_flags = ['members', 'inherited-members']
 
 
-
-
 # -- Options for HTML output ----------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
@@ -158,7 +155,7 @@ html_theme_options = {
 html_static_path = ['_static']
 
 
-# Modify the title, so as to get good social-media links
+# Modify the title to get good social-media links
 html_title = "&mdash; Dirty cat"
 
 # -- Options for HTMLHelp output ------------------------------------------
@@ -213,7 +210,7 @@ man_pages = [
 #  dir menu entry, description, category)
 texinfo_documents = [
     (master_doc, 'dirty_cat', u'dirty_cat Documentation',
-     author, 'dirty_cat', 'Learning on non-curater categorical data.',
+     author, 'dirty_cat', 'Learning on non-curated categorical data.',
      'Data Science'),
 ]
 
@@ -264,7 +261,7 @@ sphinx_gallery_conf = {
 ogp_site_url = "https://dirty-cat.github.io/stable/"
 ogp_image = "https://dirty-cat.github.io/stable/_static/dirty_cat.svg"
 ogp_use_first_image = True
-ogp_site_name = "Dirty cat"
+ogp_site_name = "dirty_cat"
 
 # -- sphinxext.napoleon configuration --------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/extensions/napoleon.html
diff --git a/doc/index.rst b/doc/index.rst
index a9452c8..96fdac9 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -49,22 +49,21 @@ pipeline you can tweak.
     </div>
     </div>
     <div class="flex-container">
-    <div class="flex-content" style="border: 0px;">
-
-For a detailed description of the problem of encoding dirty categorical data,
-see `Similarity encoding for learning with dirty categorical variables
-<https://hal.inria.fr/hal-01806175>`_ [1]_ and `Encoding high-cardinality
-string categorical variables <https://hal.inria.fr/hal-02171256v4>`_ [2]_.
 
 .. raw:: html
 
-    </div>
     <div class="flex-content" style="border: 0px;">
     <iframe style="display: block; margin: auto; width: 100%;" width="560" height="315"
      src="https://www.youtube.com/embed/_GNaaeEI2tg" frameborder="0"
      allow="accelerometer; autoplay; clipboard-write;
      encrypted-media; gyroscope; picture-in-picture" allowfullscreen
-    ></iframe>
+    ></iframe></div>
+    <div class="flex-content" style="border: 0px;">
+
+For a detailed description of the problem of encoding dirty categorical data,
+see `Similarity encoding for learning with dirty categorical variables
+<https://hal.inria.fr/hal-01806175>`_ [1]_ and `Encoding high-cardinality
+string categorical variables <https://hal.inria.fr/hal-02171256v4>`_ [2]_.
 
 .. raw:: html
 
@@ -75,7 +74,7 @@ string categorical variables <https://hal.inria.fr/hal-02171256v4>`_ [2]_.
 
    `Recent changes <CHANGES.html>`_
 
-:Installing: `$ pip install --user dirty_cat`
+:Installing: `$ pip install --user --upgrade dirty_cat`
 
 ______
 
@@ -150,7 +149,7 @@ project: there is still little known about the problems of dirty
 categories. Tradeoffs will emerge in the long run. We really need people
 giving feedback on success and failures with the different techniques and
 pointing us to open datasets on which we can do more empirical work.
-Dirty-cat received funding from `project DirtyData
+dirty-cat received funding from `project DirtyData
 <https://project.inria.fr/dirtydata/>`_ (ANR-17-CE23-0018).
 
 .. [1] Patricio Cerda, Gaël Varoquaux. Encoding high-cardinality string categorical variables. 2020. IEEE Transactions on Knowledge & Data Engineering.
diff --git a/examples/01_dirty_categories.py b/examples/01_dirty_categories.py
index 9cc11c4..89dd22d 100644
--- a/examples/01_dirty_categories.py
+++ b/examples/01_dirty_categories.py
@@ -10,7 +10,7 @@ sources.
 
 Here we look at a dataset on wages [#]_ where the column *Employee
 Position Title* contains dirty categories. On such a column, standard
-categorical encodings leads to very high dimensions and can loose
+categorical encodings leads to very high dimensions and can lose
 information on which categories are similar.
 
 We investigate various encodings of this dirty column for the machine
@@ -54,6 +54,7 @@ control.
 #
 # We first retrieve the dataset:
 from dirty_cat.datasets import fetch_employee_salaries
+
 employee_salaries = fetch_employee_salaries()
 
 # %%
@@ -69,11 +70,12 @@ y.name
 # %%
 # Now, let's carry out some basic preprocessing:
 import pandas as pd
+
 X['date_first_hired'] = pd.to_datetime(X['date_first_hired'])
 X['year_first_hired'] = X['date_first_hired'].apply(lambda x: x.year)
-# Get mask of rows with missing values in gender
+# Get a mask of the rows with missing values in "gender"
 mask = X.isna()['gender']
-# And remove the lines accordingly
+# And remove them
 X.dropna(subset=['gender'], inplace=True)
 y = y[~mask]
 
@@ -104,6 +106,7 @@ one_hot = OneHotEncoder(handle_unknown='ignore', sparse=False)
 # alongside with the column names on which each must be applied
 
 from sklearn.compose import make_column_transformer
+
 encoder = make_column_transformer(
     (one_hot, ['gender', 'department_name', 'assignment_category']),
     ('passthrough', ['year_first_hired']),
@@ -121,7 +124,9 @@ encoder = make_column_transformer(
 from sklearn.ensemble import HistGradientBoostingRegressor
 
 # We then create a pipeline chaining our encoders to a learner
+
 from sklearn.pipeline import make_pipeline
+
 pipeline = make_pipeline(encoder, HistGradientBoostingRegressor())
 
 # %%
@@ -133,15 +138,17 @@ pipeline.fit(X, y)
 # -----------------------
 #
 # The one-hot encoder is actually not well suited to the 'Employee
-# Position Title' column, as this columns contains 400 different entries:
+# Position Title' column, as this column contains 400 different entries:
 import numpy as np
+
 np.unique(y)
 
 # %%
 # We will now experiment with encoders specially made for handling
 # dirty columns
-from dirty_cat import SimilarityEncoder, TargetEncoder, MinHashEncoder,\
-    GapEncoder
+
+from dirty_cat import (SimilarityEncoder, TargetEncoder,
+                       MinHashEncoder, GapEncoder)
 
 encoders = {
     'one-hot': one_hot,
@@ -184,6 +191,7 @@ for name, method in encoders.items():
 
 import seaborn
 import matplotlib.pyplot as plt
+
 plt.figure(figsize=(4, 3))
 ax = seaborn.boxplot(data=pd.DataFrame(all_scores), orient='h')
 plt.ylabel('Encoding', size=20)
@@ -219,8 +227,9 @@ X = employee_salaries.X
 y = employee_salaries.y
 
 # %%
-# We'll drop a column we don't want
-X = X.drop(['date_first_hired'], axis=1)  # Redundant with "year_first_hired"
+# We'll drop the "date_first_hired" column as it's redundant with
+# "year_first_hired".
+X = X.drop(['date_first_hired'], axis=1)
 
 # %%
 # We still have a complex and heterogeneous dataframe:
@@ -238,7 +247,7 @@ X
 # such as gradient boosted trees, gives **a machine-learning method that
 # can be readily applied to the dataframe**.
 #
-# The SuperVectorizer requires at least dirty_cat 0.2.0. 
+# The |SV| requires at least dirty_cat 0.2.0.
 #
 
 from dirty_cat import SuperVectorizer
@@ -252,6 +261,7 @@ pipeline = make_pipeline(
 # Let's perform a cross-validation to see how well this model predicts
 
 from sklearn.model_selection import cross_val_score
+
 scores = cross_val_score(pipeline, X, y, scoring='r2')
 
 print(f'scores={scores}')
@@ -268,12 +278,13 @@ print(f'std={np.std(scores)}')
 # ------------------------------
 #
 # Let us perform the same workflow, but without the |Pipeline|, so we can
-# analyze its mechanisms along the way.
+# analyze the SuperVectorizer's mechanisms along the way.
 sup_vec = SuperVectorizer(auto_cast=True)
 
 # %%
 # We split the data between train and test, and transform them:
 from sklearn.model_selection import train_test_split
+
 X_train, X_test, y_train, y_test = train_test_split(
     X, y, test_size=0.15, random_state=42
 )
@@ -287,7 +298,7 @@ X_train_enc
 
 # %%
 # They have more columns than the original dataframe, but not much more:
-(X_train.shape, X_train_enc.shape)
+X_train.shape, X_train_enc.shape
 
 # %%
 # Inspecting the features created
@@ -295,11 +306,13 @@ X_train_enc
 #
 # The |SV| assigns a transformer for each column. We can inspect this
 # choice:
-sup_vec.transformers_
+from pprint import pprint
+
+pprint(sup_vec.transformers_)
 
 # %%
 # This is what is being passed to the |ColumnTransformer| under the hood.
-# If you're familiar with how the later works, it should be very intuitive.
+# If you're familiar with how the latter works, it should be very intuitive.
 # We can notice it classified the columns "gender" and "assignment_category"
 # as low cardinality string variables.
 # A |OneHotEncoder| will be applied to these columns.
@@ -341,12 +354,15 @@ len(feature_names)
 #    instead (which are less subject to biases)
 #
 # First, let's train the |RandomForestRegressor|,
+
 from sklearn.ensemble import RandomForestRegressor
+
 regressor = RandomForestRegressor()
 regressor.fit(X_train_enc, y_train)
 
 # %%
 # Retrieving the feature importances
+
 importances = regressor.feature_importances_
 std = np.std(
     [
@@ -363,6 +379,7 @@ indices = list(reversed(indices))
 # Plotting the results:
 
 import matplotlib.pyplot as plt
+
 plt.figure(figsize=(12, 9))
 plt.title("Feature importances")
 n = 20
diff --git a/examples/02_investigating_dirty_categories.py b/examples/02_investigating_dirty_categories.py
index 5e7314e..3ab62f6 100644
--- a/examples/02_investigating_dirty_categories.py
+++ b/examples/02_investigating_dirty_categories.py
@@ -23,30 +23,28 @@ print(data.head(n=5))
 print(data.nunique())
 
 #########################################################################
-# As we can see, some entries have many different unique values:
+# As we can see, some entries have many unique values:
 print(data['employee_position_title'].value_counts().sort_index())
 
 #########################################################################
 # These different entries are often variations on the same entities:
-# there are 3 kinds of Accountant/Auditor.
+# for example, there are 3 kinds of Accountant/Auditor.
 #
 # Such variations will break traditional categorical encoding methods:
 #
 # * Using simple one-hot encoding will create orthogonal features,
 #   whereas it is clear that those 3 terms have a lot in common.
 #
-# * If we wanted to use word embedding methods such as word2vec,
+# * If we wanted to use word embedding methods such as Word2vec,
 #   we would have to go through a cleaning phase: those algorithms
 #   are not trained to work on data such as 'Accountant/Auditor I'.
-#   However, this can be error prone and time consuming.
+#   However, this can be error-prone and time-consuming.
 #
 # The problem becomes easier if we can capture relationships between
 # entries.
 #
 # To simplify understanding, we will focus on the column describing the
 # employee's position title:
-# data
-# values = data[['employee_position_title', 'gender']] + employee_salaries.y
 
 values = data[['employee_position_title', 'gender']]
 values.insert(0, 'current_annual_salary', employee_salaries.y)
@@ -55,9 +53,10 @@ values.insert(0, 'current_annual_salary', employee_salaries.y)
 # String similarity between entries
 # -------------------------------------------------
 #
-# That's where our encoders get into play. In order to robustly
-# embed dirty semantic data, the SimilarityEncoder creates a similarity
-# matrix based on the 3-gram structure of the data.
+# That's where our encoders get into play.
+# In order to robustly embed dirty semantic data, the SimilarityEncoder
+# creates a similarity matrix based on the 3-gram structure of the data.
+
 sorted_values = values['employee_position_title'].sort_values().unique()
 
 from dirty_cat import SimilarityEncoder
@@ -70,24 +69,26 @@ transformed_values = similarity_encoder.fit_transform(
 # Plotting the new representation using multi-dimensional scaling
 # ................................................................
 #
-# Let's now plot a couple points at random using a low-dimensional representation
-# to get an intuition of what the similarity encoder is doing:
+# Let's now plot a couple of points at random using a low-dimensional
+# representation to get an intuition of what the similarity encoder is doing:
+
 from sklearn.manifold import MDS
 
 mds = MDS(dissimilarity='precomputed', n_init=10, random_state=42)
-two_dim_data = mds.fit_transform(
-    1 - transformed_values)  # transformed values lie
-# in the 0-1 range, so 1-transformed_value yields a positive dissimilarity matrix
+two_dim_data = mds.fit_transform(1 - transformed_values)
+# transformed values lie in the 0-1 range,
+# so 1-transformed_value yields a positive dissimilarity matrix
 print(two_dim_data.shape)
 print(sorted_values.shape)
 
 #########################################################################
 # We first quickly fit a KNN so that the plots does not get too busy:
+
 import numpy as np
+from sklearn.neighbors import NearestNeighbors
 
 n_points = 5
 np.random.seed(42)
-from sklearn.neighbors import NearestNeighbors
 
 random_points = np.random.choice(len(similarity_encoder.categories_[0]),
                                  n_points, replace=False)
@@ -114,6 +115,7 @@ ax.set_title(
 # ................................
 #
 # We can also plot the distance matrix for those observations:
+
 f2, ax2 = plt.subplots(figsize=(6, 6))
 cax2 = ax2.matshow(transformed_values[indices, :][:, indices])
 ax2.set_yticks(np.arange(len(indices)))
@@ -140,6 +142,7 @@ f2.tight_layout()
 #
 # A typical data-science workflow uses one-hot encoding to represent
 # categories.
+
 from sklearn.preprocessing import OneHotEncoder
 
 # encoding simply a subset of the observations
diff --git a/examples/03_feature_interpretation_gap_encoder.py b/examples/03_feature_interpretation_gap_encoder.py
index 2e7e4b4..b0cc144 100644
--- a/examples/03_feature_interpretation_gap_encoder.py
+++ b/examples/03_feature_interpretation_gap_encoder.py
@@ -15,12 +15,14 @@ categorical data.
 # --------------
 #
 # We first get the dataset:
+
 from dirty_cat.datasets import fetch_employee_salaries
 employee_salaries = fetch_employee_salaries()
 print(employee_salaries.description)
 
 ###############################################################################
 # Now, we retrieve the dirty column to encode:
+
 dirty_column = 'employee_position_title'
 X_dirty = employee_salaries.X[[dirty_column]]
 print(X_dirty.head(), end='\n\n')
@@ -31,12 +33,15 @@ print(f'Number of dirty entries = {len(X_dirty)}')
 # -------------------------
 #
 # We first create an instance of the GapEncoder with n_components=10:
+
 from dirty_cat import GapEncoder
+
 enc = GapEncoder(n_components=10, random_state=42)
 
 ###############################################################################
 # Then we fit the model on the dirty categorical data and transform it to
 # obtain encoded vectors of size 10:
+
 X_enc = enc.fit_transform(X_dirty)
 print(f'Shape of encoded vectors = {X_enc.shape}')
 
@@ -60,11 +65,12 @@ for k in range(len(topic_labels)):
 ###############################################################################
 # As expected, topics capture labels that frequently co-occur. For instance,
 # the labels *firefighter*, *rescuer*, *rescue* appear together in
-# *Firefigther/Rescuer III*, or *Fire/Rescue Lieutenant*.
+# *Firefighter/Rescuer III*, or *Fire/Rescue Lieutenant*.
 # 
 # This enables us to understand the encoding of different samples
 
 import matplotlib.pyplot as plt
+
 encoded_labels = enc.transform(X_dirty[:20])
 plt.figure(figsize=(8,10))
 plt.imshow(encoded_labels)
diff --git a/examples/04_datetime_encoder.py b/examples/04_datetime_encoder.py
index efd2c8e..e17c181 100644
--- a/examples/04_datetime_encoder.py
+++ b/examples/04_datetime_encoder.py
@@ -71,10 +71,11 @@ X_
 # The DatetimeEncoder is used by default in the SuperVectorizer, which
 # automatically detects datetime features.
 from dirty_cat import SuperVectorizer
+from pprint import pprint
 
 sup_vec = SuperVectorizer()
 sup_vec.fit_transform(X)
-sup_vec.get_feature_names_out()
+pprint(sup_vec.get_feature_names_out())
 
 ###############################################################################
 # If we want the day of the week, we can just replace SuperVectorizer's default
@@ -85,8 +86,9 @@ sup_vec.fit_transform(X)
 sup_vec.get_feature_names_out()
 
 ###############################################################################
-# We can see that the SuperVectorizer is indeed using a DatetimeEncoder for the datetime features.
-sup_vec.transformers_
+# We can see that the SuperVectorizer is indeed using
+# a DatetimeEncoder for the datetime features.
+pprint(sup_vec.transformers_)
 
 ###############################################################################
 # Predictions with date features
@@ -110,7 +112,7 @@ pipeline = make_pipeline(sup_vec, reg)
 # When using date and time features, we often care about predicting the future.
 # In this case, we have to be careful when evaluating our model, because
 # standard tools like cross-validation do not respect the time ordering.
-# Instead we can use the `TimeSeriesSplit` class, which makes sure that
+# Instead, we can use the `TimeSeriesSplit` class, which makes sure that
 # the test set is always in the future.
 X["date.utc"] = pd.to_datetime(X["date.utc"])
 sorted_indices = np.argsort(X["date.utc"])
diff --git a/examples/05_scaling_non_linear_models.py b/examples/05_scaling_non_linear_models.py
index 2f3e1ca..cb71275 100644
--- a/examples/05_scaling_non_linear_models.py
+++ b/examples/05_scaling_non_linear_models.py
@@ -2,48 +2,46 @@ r"""
 Fitting scalable, non-linear models on data with dirty categories
 =================================================================
 
-A very classic dilemma when training a machine learning model consists in
-choosing between using a linear model or a non-linear one.
+When training a machine learning model, it is often hard to choose between
+linear and a non-linear models.
 
-Linear models are very well studied and understood. They are generally fast
-and easy to optimize, and generate interpretable results. However, for some
-problems with a complex relationship between the input and the output, linear
-models reach their expressivity limit: whatever the number of samples the
-training set may have, past some point, its precision won't get any better.
+On one hand, linear models are well studied and understood. They are generally
+fast and easy to optimize, and generate interpretable results.
+They do however reach their expressiveness limits when there is a complex
+relation between the input and the output: whatever the number of samples the
+training set may have, past some point, their precision won't get any better.
 
-Non-linear models, however, tend to *scale* better with sample size: they are
-able to digest the information in the additional samples to get a better
+On the other hand, non-linear models tend to scale better with sample size:
+they are able to digest the massive amount of information and get a better
 estimate of the link between the input and the output.
 
-Non-linear models form a very large model class. Among others, this class
-includes:
+Non-linear models form a very large class. Among others, it includes:
 
 * Neural Networks
 * Tree-based methods such as Random Forests, and the very powerful Gradient
   Boosting Machines [#xgboost]_
-* Kernel Methods.
+* Kernel Methods
 
-However, reaching the phase where the non-linear model outperforms the linear
-one can be complicated. Indeed, a more complex model often means a longer
-fitting/tuning process:
+However, reaching the phase where non-linear models outperform the linear ones
+can be complicated. Indeed, a more complex model often means a longer
+fitting/tuning process
 
-* Neural networks often need extended model tuning time, in order to
-  achieve good optimization and network architecture.
-* Gradient Boosting Machines do not tend to scale extremely well with
-  increasing sample size, as all the data needs to be loaded into the main
-  memory.
+* Neural networks often need extended model tuning time in order to
+  achieve good optimization and architecture.
+* Gradient Boosting Machines tend to not scale very well with increasing
+  sample size, as all the data needs to be loaded in the main memory.
 * For kernel methods, parameter fitting requires the inversion of a gram matrix
   of size :math:`n \times n` (:math:`n` being the number of samples), yielding
-  a quadratic dependency (with :math:`n`) in the final computation time.
+  a quadratic dependency (with :math:`n`) as the final computation time.
 
 
 In order to get the best out of a non-linear model, one has to **make it
-scalable**. For kernel methods, approximation algorithms that drop the
-quadratic dependency with the sample size while ensuring almost the
-same model capacity exist.
+scalable**. Quite notably, for kernel methods, approximation algorithms that
+drop the quadratic dependency with the sample size while ensuring capacity
+approximation exist.
 
 In this example, you will learn how to:
-    1. Build a ML pipeline that uses a kernel method.
+    1. Build a machine learning pipeline that uses a kernel method.
     2. Make this pipeline scalable, by using online algorithms and dimension
        reduction methods.
 
@@ -61,52 +59,40 @@ In this example, you will learn how to:
 .. |NYS_EXAMPLE|
     replace:: :ref:`scikit-learn documentation <nystroem_kernel_approx>`
 
-.. |RBF|
-    replace:: :class:`~sklearn.kernel_approximation.RBFSampler`
+.. |RBF| replace:: :class:`~sklearn.kernel_approximation.RBFSampler`
 
-.. |SVC|
-    replace:: :class:`SupportVectorClassifier <sklearn.svm.SVC>`
+.. |SVC| replace:: :class:`SupportVectorClassifier <sklearn.svm.SVC>`
 
 .. |SE| replace:: :class:`~dirty_cat.SimilarityEncoder`
 
-.. |SGDClassifier| replace::
-    :class:`~sklearn.linear_model.SGDClassifier`
+.. |SGDClassifier| replace:: :class:`~sklearn.linear_model.SGDClassifier`
 
-.. |APS| replace::
-    :func:`~sklearn.metrics.average_precision_score`
+.. |APS| replace:: :func:`~sklearn.metrics.average_precision_score`
 
-.. |OneHotEncoder| replace::
-    :class:`~sklearn.preprocessing.OneHotEncoder`
+.. |OHE| replace:: :class:`~sklearn.preprocessing.OneHotEncoder`
 
-.. |ColumnTransformer| replace::
-    :class:`~sklearn.compose.ColumnTransformer`
+.. |ColumnTransformer| replace:: :class:`~sklearn.compose.ColumnTransformer`
 
-.. |LabelEncoder| replace::
-    :class:`~sklearn.preprocessing.LabelEncoder`
+.. |LabelEncoder| replace:: :class:`~sklearn.preprocessing.LabelEncoder`
 
 .. |SGDClassifier_partialfit| replace::
     :meth:`~sklearn.linear_model.SGDClassifier.partial_fit`
 
-.. |Pipeline| replace::
-    :class:`~sklearn.pipeline.Pipeline`
+.. |Pipeline| replace:: :class:`~sklearn.pipeline.Pipeline`
 
-.. |pd_read_csv| replace::
-    :func:`pandas.read_csv`
+.. |pd_read_csv| replace:: :func:`pandas.read_csv`
 
-.. |sklearn| replace::
-    :std:doc:`scikit-learn <sklearn:index>`
+.. |sklearn| replace:: :std:doc:`scikit-learn <sklearn:index>`
 
-.. |transform| replace::
-    :std:term:`transform <sklearn:transform>`
+.. |transform| replace:: :std:term:`transform <sklearn:transform>`
 
-.. |fit| replace::
-    :std:term:`fit <sklearn:fit>`
+.. |fit| replace:: :std:term:`fit <sklearn:fit>`
 
 """
 ###############################################################################
 # Training a simple pipeline
 # --------------------------
-# The data that the model will fit is the :code:`drug_directory` dataset.
+# We will train our model on the :code:`drug_directory` dataset.
 from dirty_cat.datasets import fetch_drug_directory
 
 # We'll only gather the information, and not load the dataset in memory for now
@@ -114,25 +100,23 @@ drug_directory = fetch_drug_directory(load_dataframe=False)
 print(drug_directory.description)
 
 ###############################################################################
-# .. topic:: Problem Setting
-#
-#    We set the goal of our machine learning problem as follows:
+# .. topic:: Problem Setting:
 #
 #    .. centered::
-#       **predict the type of a drug given its composition.**
+#       **predict the drug type given its composition**
 #
 #
-# The :code:`NONPROPRIETARYNAME` column, is composed of text observations with
+# The :code:`NONPROPRIETARYNAME` column, is composed of textual observations
 # describing each drug's composition. The :code:`PRODUCTTYPENAME` column
-# consists of categorical values: therefore, our problem is a classification
-# problem. You can have a glimpse of the values here:
+# consists of categorical values; we therefore a classification problem.
+#
+# You can have a glimpse of the values here:
 import pandas as pd
 
 df = pd.read_csv(
     drug_directory.path,
-    quotechar="'",
-    escapechar='\\',
     nrows=10,
+    **drug_directory.read_csv_kwargs,
 ).astype(str)
 print(df[['NONPROPRIETARYNAME', 'PRODUCTTYPENAME']].head())
 # This will be useful further down in the example.
@@ -141,33 +125,34 @@ columns_names = df.columns
 ###############################################################################
 # Estimators construction
 # -----------------------
-# Our input is categorical, thus needs to be encoded. As observations often
-# consist in variations around a few concepts (for instance,
-# :code:`'Amlodipine Besylate'` and
-# :code:`'Amlodipine besylate and atorvastatin calcium'`
-# have one ingredient in common), we need an encoding able to
-# capture similarities between observations.
-
+# Our input is categorical, thus needs to be encoded.
+#
+# As observations often consist in variations around a few concepts
+# - for instance, "Amlodipine Besylate" and
+# "Amlodipine besylate and atorvastatin calcium"
+# have one ingredient in common - we need an encoding able to
+# capture similarities between observations. We choose to use a |SE|.
+#
+# Two other columns are used to predict the output: ``DOSAGEFORMNAME`` and
+# ``ROUTENAME``.
+# They are both categorical and can be encoded with a |OHE|.
+# We use a |ColumnTransformer| to stack both.
+# We can now choose a kernel method. Here, we'll use a |SVC| to fit
+# the encoded inputs.
 from dirty_cat import SimilarityEncoder
-from sklearn.compose import make_column_transformer
 from sklearn.preprocessing import OneHotEncoder
-similarity_encoder = SimilarityEncoder()
-
-###############################################################################
-# Two other columns are used to predict the output: ``DOSAGEFORMNAME`` and
-# ``ROUTENAME``. They are both categorical and can be encoded with a
-# |OneHotEncoder|. We use a |ColumnTransformer| to stack the |OneHotEncoder|
-# and the |SE|.  We can now choose a kernel method, for instance a |SVC|, to
-# fit the encoded inputs.
-from sklearn.pipeline import Pipeline
-from sklearn.svm import SVC
+from sklearn.compose import make_column_transformer
 
 column_transformer = make_column_transformer(
-    (similarity_encoder, ['NONPROPRIETARYNAME']),
+    (SimilarityEncoder(), ['NONPROPRIETARYNAME']),
     (OneHotEncoder(handle_unknown='ignore'), ['DOSAGEFORMNAME', 'ROUTENAME']),
-    sparse_threshold=1)
+    sparse_threshold=1,
+)
 
-# The classifier and the ColumnTransformer are stacked into a Pipeline object
+###############################################################################
+# We will bundle the ColumnTransformer and the classifier into a |Pipeline|
+from sklearn.svm import SVC
+from sklearn.pipeline import Pipeline
 classifier = SVC(kernel='rbf', random_state=42, gamma=1)
 steps = [('transformer', column_transformer), ('classifier', classifier)]
 model = Pipeline(steps)
@@ -175,29 +160,52 @@ model = Pipeline(steps)
 ###############################################################################
 # Data Loading and Preprocessing
 # ------------------------------
-# Like in most machine learning setups, the data has to be split into 2
+#
+# As usual in most machine learning setups, the data has to be split into 2
 # exclusive parts:
 #
-# * One for model training.
-# * One for model testing.
+# As usual in machine learning settings, we will split our data in two
+# exclusive sets:
 #
-# For this reason, we create a simple wrapper around |pd_read_csv|, that
-# extracts the :code:`X`, and :code:`y` from the dataset.
+# * One for model training
+# * One for model testing
+#
+# We create a simple wrapper around |pd_read_csv|, that extracts :code:`X`,
+# and :code:`y` from the dataset.
 #
 # .. topic:: Note about class imbalance:
 #
-#    The :code:`y` labels are composed of 7 unique classes. However, ``HUMAN
-#    OTC DRUG`` and ``HUMAN PRESCRIPTION DRUG`` represent around 97% of the
-#    data, in a fairly balanced manner. The last 5 classes are much rarer.
+#    The :code:`y` labels are composed of 7 unique classes. However, the two
+#    classes ``HUMAN OTC DRUG`` and ``HUMAN PRESCRIPTION DRUG`` represent
+#    around 97% of the data. The other 5 classes are much rarer.
 #    Dealing with class imbalance is out of the scope of this example, so the
-#    models will be trained on the first two classes only.
+#    models will be trained only on the two most common classes.
+#
+# Classifiers in |sklearn| often require :code:`y` to be integer labels.
+# Additionally, |APS| requires a binary version of the labels.
+#
+# We therefore create for these respective constraints a |LabelEncoder| and
+# a |OHE|
+import numpy as np
+from typing import Tuple
+from sklearn.preprocessing import LabelEncoder, OneHotEncoder
+
 
+label_encoder = LabelEncoder()
+# Pre-fit on the known classes of y
+label_encoder.fit(['HUMAN OTC DRUG', 'HUMAN PRESCRIPTION DRUG'])
 
-def preprocess(df, label_encoder):
+one_hot_encoder = OneHotEncoder(categories="auto", sparse=False)
+# Pre-fit on the integer labels
+one_hot_encoder.fit([[0], [1]])
+
+
+def preprocess(df: pd.DataFrame) -> Tuple[pd.DataFrame, np.array]:
     df = df.loc[df['PRODUCTTYPENAME'].isin(
         ['HUMAN OTC DRUG', 'HUMAN PRESCRIPTION DRUG'])]
 
-    df = df[['NONPROPRIETARYNAME', 'DOSAGEFORMNAME', 'ROUTENAME', 'PRODUCTTYPENAME']]
+    df = df[['NONPROPRIETARYNAME', 'DOSAGEFORMNAME',
+             'ROUTENAME', 'PRODUCTTYPENAME']]
     df = df.dropna()
 
     X = df[['NONPROPRIETARYNAME', 'DOSAGEFORMNAME', 'ROUTENAME']]
@@ -208,79 +216,57 @@ def preprocess(df, label_encoder):
     return X, y_int
 
 
-def get_X_y(**kwargs):
-    """simple wrapper around pd.read_csv that extracts features and labels
+def get_X_y(**kwargs) -> Tuple[pd.DataFrame, np.array]:
+    """Simple wrapper around pd.read_csv that extracts features and labels.
 
     Some systematic preprocessing is also carried out to avoid doing this
     transformation repeatedly in the code.
     """
-    global label_encoder
-    df = pd.read_csv(drug_directory.path, **drug_directory.read_csv_kwargs, **kwargs)
-    return preprocess(df, label_encoder)
+    df = pd.read_csv(drug_directory.path, names=columns_names,
+                     **drug_directory.read_csv_kwargs, **kwargs)
+    return preprocess(df)
 
-###############################################################################
-# Classifier objects in |sklearn| often require :code:`y` to be integer labels.
-# Additionally, |APS| requires a binary version of the labels.  For these two
-# purposes, we create:
-#
-# * a |LabelEncoder|, that we pre-fitted on the known :code:`y` classes
-# * a |OneHotEncoder|, pre-fitted on the resulting integer labels.
-#
-# Their |transform| methods can the be called at appropriate times.
-import numpy as np
-from sklearn.preprocessing import LabelEncoder, OneHotEncoder
-
-label_encoder = LabelEncoder()
-label_encoder.fit(['HUMAN OTC DRUG', 'HUMAN PRESCRIPTION DRUG'])
-
-one_hot_encoder = OneHotEncoder(categories="auto", sparse=False)
-one_hot_encoder.fit([[0], [1]])
 
 ###############################################################################
 # .. warning::
 #
-#    During the following training procedures of this example, we will assume
-#    that the dataset was shuffled prior to its loading. As a reason, we can
-#    take the first :math:`n` observations for the training set, the next
-#    :math:`m` observations for the test set and so on. This may not be the
-#    case for all datasets, so be careful before applying this code to your own
-#    !
+#    In the following steps, we assume the dataset to be shuffled prior to its
+#    loading. Be careful as this is not the case for all datasets!
+#    This is important here because we will take the :math:`n` first
+#    observations for the training set and the next :math:`m` observations for
+#    the test set repeatably.
 #
 #
-# Finally, the :code:`X` and :code:`y` are loaded.
+# Finally, :code:`X` and :code:`y` are loaded.
 #
 #
 # .. topic:: Note: offsetting the test set
 #
-#    We create an offset to separate the training and the test set. The reason
-#    for this, is that the online procedures of this example will consume far
-#    more rows, but we still would like to compare accuracies with the
-#    same test set, and not change it each time. Therefore, we "reserve" the
-#    first 100000 rows for the training phase. The rest is made available to
-#    the test set.
+#    For the online-learning procedures of this example, we would like to
+#    compare the accuracy at each iteration. In order to do that, we need to
+#    keep the same test set. As such, we "reserve" the first 100k rows for
+#    the training phase, and use the rest as the test set.
 train_set_size = 5000
 test_set_size = 10000
 offset = 100000
 
-X_train, y_train = get_X_y(skiprows=1, names=columns_names,
-                           nrows=train_set_size)
+X_train, y_train = get_X_y(skiprows=1, nrows=train_set_size)
 
-X_test, y_test = get_X_y(skiprows=offset, names=columns_names,
-                         nrows=test_set_size)
+X_test, y_test = get_X_y(skiprows=offset, nrows=test_set_size)
 
 ###############################################################################
 # Evaluating time and sample complexity
 # -------------------------------------
+#
 # Let's get an idea of model precision and performance depending on the number
 # of the samples used in the train set.
 # The |Pipeline| is trained over different training set sizes. For this,
 # :code:`X_train` and :code:`y_train` get sliced into subsets of increasing
-# size, while :code:`X_test` and :code:`y_test` do not change when the
-# sample size varies.
+# size, while :code:`X_test` and :code:`y_test` do not change.
 import time
 from sklearn.metrics import average_precision_score
 
-# define the different train set sizes on which to evaluate the model
+# Define the different train set sizes
 train_set_sizes = [train_set_size // 10, train_set_size // 3, train_set_size]
 
 train_times_svc, test_scores_svc = [], []
@@ -301,27 +287,29 @@ for n in train_set_sizes:
     train_times_svc.append(train_time)
     test_scores_svc.append(test_score)
 
-    msg = ("using {:>5} samples: model fitting took {:.1f}s, test accuracy of "
-           "{:.3f}")
-    print(msg.format(n, train_time, test_score))
+    print(
+        f"Using {n:>5} samples: model fitting took {train_time:.1f}s, "
+        f"test accuracy of {test_score:.3f}. "
+    )
 
 
 ###############################################################################
-# Increasing training size clearly improves model accuracy. However, the
-# training time and the input size increase quadratically with the training set
-# size.  Indeed, kernel methods need to process an entire :math:`n \times n`
+# Increasing the training set size clearly improves model accuracy.
+# However, the training time and the input size increase quadratically with it.
+#
+# Indeed, kernel methods need to process an entire :math:`n \times n`
 # matrix at once. In order for this matrix to be loaded into memory, :math:`n`
-# has to remain low: Using 30000 observations, the input is a :math:`30000
-# \times 30000` matrix.  If composed of 32-bit floats, its total size is around
+# has to remain low: If we used 32-bit floats, and with 30k observations, the
+# total size of this 30k x 30k matrix would be around
 # :math:`30000^2 \times 32 = 2.8 \times 10^{10} \text{bits} = 4\text{GB}`
 
 ###############################################################################
-# Reducing input dimension using kernel approximation methods.
-# ------------------------------------------------------------
+# Reducing input dimension using kernel approximation methods
+# -----------------------------------------------------------
 #
 # The main scalability issues with kernels methods is the processing of a large
 # square matrix. To understand where this matrix comes from, we need to delve a
-# little bit deeper these methods internals.
+# bit deeper into these methods.
 #
 # .. topic:: Kernel methods
 #
@@ -329,41 +317,45 @@ for n in train_set_sizes:
 #    between each pair of inputs. Using a similarity matrix to solve a machine
 #    learning problem allows to catch complex, non-linear relationships within
 #    the data.  But it requires inverting this matrix, which can be a
-#    computational burden when the sample sizes increases.
+#    computational burden when the sample size increases.
 
 
 ###############################################################################
 # Kernel approximation methods
 # ----------------------------
-# From what was said below, two criteria are limiting a kernel algorithm to
-# scale:
+#
+# From what we saw above, two criterion are limiting kernel algorithms from
+# scaling:
 #
 # * It processes a matrix, whose size increases quadratically with the number
 #   of samples.
-# * During fitting time, this matrix is **inverted**, meaning it has to be
+# * During training, this matrix is **inverted**, meaning it has to be
 #   loaded into main memory.
 #
 # Kernel approximation methods such as |RBF| or |NYS| [#nys_ref]_ try to
 # approximate this similarity matrix, without actually creating it. By allowing
-# the program to not compute the perfect similarity matrix, the problem
-# complexity becomes linear! Plus, the samples also do not need to be processed
-# at once into main memory. We are not bound to use a |SVC| anymore, and can
-# instead use an online optimization that will process the input by batch.
+# the program to not compute the perfect similarity matrix, **the problem
+# complexity becomes linear**! Furthermore, the samples don't need to be
+# processed at once.
+#
+# We are not bound to use a |SVC| anymore, and can instead use a |RBF| in an
+# online optimization setting to process the input by batch.
 #
 # .. topic:: Online algorithms
 #
 #    An online algorithm [#online_ref]_ is an algorithm that treats its input
-#    piece by piece in a serial fashion. An famous example is the stochastic
-#    gradient descent [#sgd_ref]_, where an estimation the objective function's
-#    gradient is computed on a batch of the data at each step.
-#
+#    piece by piece in a serial fashion.
+#    A famous example is the stochastic gradient descent [#sgd_ref]_,
+#    where an estimation of the objective function's gradient is computed
+#    on a batch of the data at each step.
 #
 
 ###############################################################################
 # Reducing the transformers dimensionality
 # ----------------------------------------
+#
 # There is one last scalability issue in our pipeline: the |SE| and the |RBF|
-# both implement the |fit| method. How to adapt those method to an online
+# both implement the |fit| method. How to adapt those methods to an online
 # setting, where the data is never loaded as a whole?
 #
 # A simple solution is to partially fit the |SE| and the |RBF| on a subset of
@@ -374,16 +366,19 @@ n_out_encoder = 1000
 n_out_rbf = 5000
 n_samples_encoder = 10000
 
-X_encoder, _ = get_X_y(nrows=n_samples_encoder, names=columns_names)
-
-similarity_encoder = SimilarityEncoder(
-    categories='most_frequent', n_prototypes=n_out_encoder,
-    random_state=42, ngram_range=(2, 4))
+X_encoder, _ = get_X_y(nrows=n_samples_encoder)
 
 # Fit the rbf_sampler with the similarity matrix.
 column_transformer = make_column_transformer(
-    (similarity_encoder, ['NONPROPRIETARYNAME']),
-    (OneHotEncoder(handle_unknown='ignore'), ['DOSAGEFORMNAME', 'ROUTENAME']),
+    (
+        SimilarityEncoder(categories='most_frequent',
+                          n_prototypes=n_out_encoder,
+                          random_state=42, ngram_range=(2, 4)),
+        ['NONPROPRIETARYNAME']
+    ), (
+        OneHotEncoder(handle_unknown='ignore'),
+        ['DOSAGEFORMNAME', 'ROUTENAME']
+    ),
     sparse_threshold=1)
 
 transformed_categories = column_transformer.fit_transform(X_encoder)
@@ -416,23 +411,22 @@ X_test_kernel_approx, y_true_test_onehot = encode(
 ###############################################################################
 # Online training for out-of-memory data
 # --------------------------------------
-# We now have all the theoretical elements to create an non-linear, online
-# kernel method.
+# We now have all the elements to create a non-linear, online kernel method.
 import warnings
 from sklearn.linear_model import SGDClassifier
 
 online_train_set_size = 100000
-# Filter warning on max_iter and tol
+# Filter warnings on max_iter and tol
 warnings.filterwarnings('ignore', module='sklearn.linear_model')
 sgd_classifier = SGDClassifier(
     max_iter=1, tol=None, random_state=42, average=10)
 
 
 ###############################################################################
-# We can now start the training, by looping over batches one by one. Note that
-# only one pass over the whole dataset is done. It may be worth doing several
-# passes, but for very large sample sizes, the increase in test accuracy is
-# likely to be marginal.
+# We can now start the training by looping over batches one by one.
+# Note that only one pass over the whole dataset is done.
+# It may be worth doing several passes, but for very large sample sizes,
+# the increase in test accuracy is likely to be marginal.
 batchsize = 1000
 test_scores_rbf = []
 train_times_rbf = []
@@ -444,20 +438,20 @@ iter_csv = pd.read_csv(
     skiprows=1, names=columns_names, **drug_directory.read_csv_kwargs)
 
 for batch_no, batch in enumerate(iter_csv):
-    X_batch, y_batch = preprocess(batch, label_encoder)
-    # skip iteration if batch is empty after preprocessing
+    X_batch, y_batch = preprocess(batch)
+    # Skip iteration if batch is empty after preprocessing
     if len(y_batch) == 0:
         continue
     X_batch_kernel_approx, y_batch_onehot = encode(
         X_batch, y_batch, one_hot_encoder, column_transformer, rbf_sampler)
 
-    # make one pass of stochastic gradient descent over the batch.
+    # Make one pass of stochastic gradient descent over the batch.
     sgd_classifier.partial_fit(
         X_batch_kernel_approx, y_batch, classes=[0, 1])
 
-    # print train/test accuracy metrics every 5 batch
+    # Print train/test accuracy metrics every 5 batch
     if (batch_no % 5) == 0:
-        message = "batch {:>4} ".format(batch_no)
+        message = f"batch {batch_no:>4} "
         for origin, X, y_true_onehot in zip(
                 ('train', 'val'),
                 (X_batch_kernel_approx, X_test_kernel_approx),
@@ -465,24 +459,26 @@ for batch_no, batch in enumerate(iter_csv):
 
             y_pred = sgd_classifier.predict(X)
 
-            # preprocess correctly the labels and prediction to match
+            # Preprocess correctly the labels and prediction to match
             # average_precision_score expectations
             y_pred_onehot = one_hot_encoder.transform(
                 y_pred.reshape(-1, 1))
 
             score = average_precision_score(y_true_onehot, y_pred_onehot)
-            message += "{} precision: {:.4f}  ".format(origin, score)
+            message += f"{origin} precision: {score:.4f}  "
             if origin == 'val':
                 test_scores_rbf.append(score)
                 train_times_rbf.append(time.perf_counter() - t0)
-                online_train_set_sizes.append((batch_no + 1)*batchsize)
+                online_train_set_sizes.append((batch_no + 1) * batchsize)
 
         print(message)
 
 ###############################################################################
-# So far, we fitted two kinds of models: a exact kernel algorithm, and an
-# approximate, online one. Lets compare both the accuracies and the number of
-# visited samples for each model as we increase our time budget:
+# So far, we fitted two kinds of models: an exact kernel algorithm, and an
+# approximate, online one.
+#
+# Let's compare both the accuracies and the number of visited samples
+# for each model as we increase our time budget:
 import matplotlib.pyplot as plt
 
 f, axs = plt.subplots(2, 1, sharex=True, figsize=(10, 7))
@@ -499,17 +495,17 @@ ax_capacity.plot(train_times_svc, train_set_sizes, 'b-', label='exact')
 ax_capacity.plot(train_times_rbf, online_train_set_sizes, 'r-', label='online')
 ax_capacity.set_yscale('log')
 
-ax_score.legend(
-    bbox_to_anchor=(0., 1.02, 1., .102),
-    loc=3, ncol=2, mode='expand',
-    borderaxespad=0.)
+ax_score.legend(bbox_to_anchor=(0., 1.02, 1., .102),
+                loc=3, ncol=2, mode='expand',
+                borderaxespad=0.)
 
-# compare the two methods in their common time range
+# Compare the two methods in the common time range
 ax_score.set_xlim(0, min(train_times_svc[-1], train_times_rbf[-1]))
 
-title = """Test set accuracy and number of samples visited
-samples seen by the SGDClassifier"""
-f.suptitle(title)
+f.suptitle("Test set accuracy and number of samples visited"
+           "samples seen by the SGDClassifier")
+
+f.show()
 
 
 ###############################################################################
@@ -518,15 +514,20 @@ f.suptitle(title)
 # non-online, exact kernel method.
 #
 # Our online model also outperforms online **linear** models (for instance,
-# |SE| + |SGDClassifier|). We did not fit two online models here for simplicity
-# purposes, but to train an online linear model, simply comment out the line in
-# encode :code:`X_highdim = rbf_sampler.transform(X_sim_encoded.toarray())` and
-# change it for example with :code:`X_highdim = X_sim_encoded`.
+# |SE| + |SGDClassifier|).
+#
+# .. note::
+#
+#    We did not fit two online models here for the sake of simplicity, but to
+#    train an online linear model, simply comment out the line
+#    :code:`X_highdim = rbf_sampler.transform(X_sim_encoded.toarray())`
+#    in the function :code:`encode` and change it with, for example,
+#    :code:`X_highdim = X_sim_encoded`.
 #
-# In particular, this hierarchy between the linear model and the non-linear
-# one shows that there were some significant non-linear relationships
+# In particular, this hierarchy between the linear and the non-linear models
+# shows that there were some significant non-linear relationships
 # between the input and the output. By scaling a kernel method, we successfully
-# took this non-linearity into account in our model, which was a far from
+# took this non-linearity into account in our model, which was far from a
 # trivial task at the beginning of this example!
 
 ###############################################################################
diff --git a/examples/06_dimension_reduction_and_performance.py b/examples/06_dimension_reduction_and_performance.py
index 32dd591..f37c069 100644
--- a/examples/06_dimension_reduction_and_performance.py
+++ b/examples/06_dimension_reduction_and_performance.py
@@ -1,20 +1,27 @@
 """
 Scalability considerations for similarity encoding
-===================================================
+==================================================
 
-Here we discuss how to apply efficiently SimilarityEncoder to larger
+We discuss in this notebook how to efficiently apply the |SE| to larger
 datasets: reducing the number of reference categories to "prototypes",
 either chosen as the most frequent categories, or with kmeans clustering.
 
-Note that the :class:`GapEncoder` naturally does data reduction and comes
-with online estimation. As a result is it more scalable than the
-SimilarityEncoder, and should be preferred in large-scale settings.
+
+.. note::
+    The |Gap| naturally does data reduction and comes with online estimation.
+    As a result, is it more scalable than the |SE|,
+    and should be preferred in large-scale settings.
+
+
+.. |SE| replace:: :class:`~dirty_cat.SimilarityEncoder`
+
+.. |Gap| replace:: :class:`~dirty_cat.GapEncoder`
+
+.. |ColumnTransformer| replace:: :class:`~sklearn.compose.ColumnTransformer`
+
+.. |OHE| replace:: :class:`~sklearn.preprocessing.OneHotEncoder`
 
 """
-# Avoid the warning in scikit-learn's LogisticRegression for the change
-# in the solver
-import warnings
-warnings.simplefilter(action='ignore', category=FutureWarning)
 
 ###############################################################################
 # A tool to report memory usage and run time
@@ -22,25 +29,26 @@ warnings.simplefilter(action='ignore', category=FutureWarning)
 #
 # For this example, we build a small tool that reports memory
 # usage and compute time of a function
-from time import time
+from time import perf_counter
 import functools
 import tracemalloc
 
 
 def resource_used(func):
-    """ Decorator that return a function that prints its usage
+    """
+    Decorator for performance analysis.
     """
 
     @functools.wraps(func)
     def wrapped_func(*args, **kwargs):
-        t0 = time()
+        t0 = perf_counter()  # Launch a time
         tracemalloc.start()
         out = func(*args, **kwargs)
         size, peak = tracemalloc.get_traced_memory()
         tracemalloc.stop()
         peak /= (1024 ** 2)  # Convert to megabytes
-        print("Run time: %.1is    Memory used: %iMb"
-              % (time() - t0, peak))
+        print(f"Run time: {perf_counter() - t0:.2f}s | "
+              f"Memory used: {peak:.2f}MB. ")
         return out
 
     return wrapped_func
@@ -50,29 +58,31 @@ def resource_used(func):
 # Data Importing and preprocessing
 # --------------------------------
 #
-# We first get the dataset:
+# First, let's fetch the dataset we'll use further down
 import pandas as pd
 from dirty_cat.datasets import fetch_open_payments
 
 open_payments = fetch_open_payments()
-print(open_payments.description)
+X = open_payments.X
 
-df = open_payments.X
-
-na_mask: pd.DataFrame = df.isna()
-df = df.dropna(axis=0)
-df = df.reset_index()
+open_payments.description
 
+###############################################################################
+# We'll perform a some cleaning
 from functools import reduce
 
+# Remove the missing lines in X
+na_mask: pd.DataFrame = X.isna()
+X = X.dropna(axis=0).reset_index(drop=True)
+
 y = open_payments.y
 # Combine boolean masks
-na_mask = reduce(lambda acc, col: acc | na_mask[col],
-                 na_mask.columns, na_mask[na_mask.columns[0]])
-# Drop the lines that contained missing values in X
-y = y[~na_mask]
-y.reset_index()
+na_mask = na_mask.any(axis=1)
+# Drop the lines in y that contained missing values in X
+y = y[~na_mask].reset_index(drop=True)
 
+###############################################################################
+# We'll write down which columns are clean and which are dirty
 clean_columns = [
     'Applicable_Manufacturer_or_Applicable_GPO_Making_Payment_Name',
     'Dispute_Status_for_Publication',
@@ -84,52 +94,54 @@ dirty_columns = [
 ]
 
 ###############################################################################
-# We will use SimilarityEncoder on the the two dirty columns defined above.
-# One difficulty is that they have many different entries.
-print(df[dirty_columns].nunique())
-
-###############################################################################
-print(df[dirty_columns].value_counts()[:20])
+# We will use |SE| on the two dirty columns defined above.
+# One difficulty is that they have many entries, and because of that, as we'll
+# see, the :code:`SimilarityEncoder` will take a while.
+X[dirty_columns].value_counts()[:20]
 
 ###############################################################################
-# As we will see, SimilarityEncoder takes a while on such data.
-
+X[dirty_columns].nunique()
 
 ###############################################################################
 # SimilarityEncoder with default options
 # --------------------------------------
 #
-# Let us build our vectorizer, using a ColumnTransformer to combine
-# one-hot encoding and similarity encoding
+# Let us build our vectorizer, using a |ColumnTransformer| to combine
+# a |OHE| and a |SE|
 from sklearn.preprocessing import OneHotEncoder
 from sklearn.compose import ColumnTransformer
 from dirty_cat import SimilarityEncoder
 
-sim_enc = SimilarityEncoder()
-
-transformers = [
-    ('one_hot', OneHotEncoder(sparse=False, handle_unknown='ignore'), clean_columns),
+clean_col_transformer = [
+    ('one_hot',
+     OneHotEncoder(sparse=False, handle_unknown='ignore'),
+     clean_columns),
 ]
 
 column_trans = ColumnTransformer(
-    transformers=transformers + [('sim_enc', sim_enc, dirty_columns)],
+    transformers=clean_col_transformer + [
+        ('sim_enc',
+         SimilarityEncoder(),
+         dirty_columns)
+    ],
     remainder='drop')
 
-t0 = time()
-X = column_trans.fit_transform(df)
-t1 = time()
-print('Time to vectorize: %s' % (t1 - t0))
+t0 = perf_counter()
+X_enc = column_trans.fit_transform(X)
+t1 = perf_counter()
+print(f'Time to vectorize: {t1 - t0:.3f}s')
 
 ###############################################################################
-# We can run a cross-validation
-from sklearn import linear_model, pipeline, model_selection
+# Let's now run a cross-validation!
+from sklearn import pipeline, model_selection
+from sklearn.linear_model import LogisticRegression
 
 # We specify max_iter to avoid convergence warnings
-log_reg = linear_model.LogisticRegression(max_iter=10000)
+log_reg = LogisticRegression(max_iter=10000)
 
 model = pipeline.make_pipeline(column_trans, log_reg)
-results = resource_used(model_selection.cross_validate)(model, df, y)
-print("Cross-validation score: %s" % results['test_score'])
+results = resource_used(model_selection.cross_validate)(model, X, y)
+print(f"Cross-validation score: {results['test_score']}")
 
 ###############################################################################
 # Store results for later
@@ -142,20 +154,24 @@ times['Default options'] = results['fit_time']
 # Most frequent strategy to define prototypes
 # -------------------------------------------
 #
-# The most frequent strategy selects the n most frequent values in a dirty
-# categorical variable to reduce the dimensionality of the problem and thus
-# speed things up. We select manually the number of prototypes we want to use.
-sim_enc = SimilarityEncoder(categories='most_frequent', n_prototypes=100)
+# The :code:`most_frequent` strategy selects the :code:`n` most frequent
+# values in a dirty categorical variable to reduce the dimensionality of the
+# problem and thus speed things up.
+# Here, we arbitrarily choose 100 as the number of prototypes we want to use.
 
 column_trans = ColumnTransformer(
-    transformers=transformers + [('sim_enc', sim_enc, dirty_columns)],
+    transformers=clean_col_transformer + [
+        ('sim_enc',
+         SimilarityEncoder(categories='most_frequent', n_prototypes=100),
+         dirty_columns)
+    ],
     remainder='drop')
 
 ###############################################################################
-# Check now that prediction is still as good
+# Check that the prediction is still as good
 model = pipeline.make_pipeline(column_trans, log_reg)
-results = resource_used(model_selection.cross_validate)(model, df, y)
-print("Cross-validation score: %s" % results['test_score'])
+results = resource_used(model_selection.cross_validate)(model, X, y)
+print(f"Cross-validation score: {results['test_score']}")
 
 ###############################################################################
 # Store results for later
@@ -166,19 +182,23 @@ times['Most frequent'] = results['fit_time']
 # KMeans strategy to define prototypes
 # ------------------------------------
 #
-# K-means strategy is also a dimensionality reduction technique.
-# SimilarityEncoder can apply a K-means and nearest neighbors algorithm
-# to find the prototypes. The number of prototypes is set manually.
-sim_enc = SimilarityEncoder(categories='k-means', n_prototypes=100)
+# The k-means strategy is also a dimensionality reduction technique.
+# The :code:`SimilarityEncoder` can apply a K-means and nearest neighbors
+# algorithm to find the prototypes. Once again, the number of prototypes
+# we chose here is arbitrary.
 
 column_trans = ColumnTransformer(
-    transformers=transformers + [('sim_enc', sim_enc, dirty_columns)],
+    transformers=clean_col_transformer + [
+        ('sim_enc',
+         SimilarityEncoder(categories='k-means', n_prototypes=100),
+         dirty_columns)
+    ],
     remainder='drop')
 
 ###############################################################################
-# Check now that prediction is still as good
+# Check that the prediction is still as good
 model = pipeline.make_pipeline(column_trans, log_reg)
-results = resource_used(model_selection.cross_validate)(model, df, y)
+results = resource_used(model_selection.cross_validate)(model, X, y)
 print("Cross-validation score: %s" % results['test_score'])
 
 ###############################################################################
@@ -187,8 +207,8 @@ scores['KMeans'] = results['test_score']
 times['KMeans'] = results['fit_time']
 
 ###############################################################################
-# Plot a summary figure
-# ---------------------
+# Summary
+# -------
 import seaborn
 import matplotlib.pyplot as plt
 
diff --git a/examples/README.txt b/examples/README.txt
index c4ed750..b3fafcf 100644
--- a/examples/README.txt
+++ b/examples/README.txt
@@ -1,3 +1,3 @@
 
 Using dirty_cat
-===============
+===============
\ No newline at end of file
diff --git a/requirements-min.txt b/requirements-min.txt
index 4ae6b4c..05ed775 100644
--- a/requirements-min.txt
+++ b/requirements-min.txt
@@ -9,4 +9,3 @@ pandas==1.1.5
 cython
 joblib
 seaborn
-
diff --git a/setup.py b/setup.py
index 173e1b3..6c85090 100755
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,7 @@ if __name__ == '__main__':
           description=("Machine learning with dirty categories."),
           long_description=DESCRIPTION,
           license='BSD',
+          python_requires='>=3.6',
           classifiers=[
               'Development Status :: 5 - Production/Stable',
               'Environment :: Console',


Test patch:
diff --git a/dirty_cat/datasets/tests/test_fetching.py b/dirty_cat/datasets/tests/test_fetching.py
index 24fa6af..fbe25ee 100644
--- a/dirty_cat/datasets/tests/test_fetching.py
+++ b/dirty_cat/datasets/tests/test_fetching.py
@@ -1,16 +1,9 @@
 """
-
 Tests fetching.py (datasets fetching off OpenML.org).
-
 """
 
-# Author:
-# Lilian Boulard <lilian@boulard.fr>
-# https://github.com/LilianBoulard
-
 import pytest
 import shutil
-import sklearn
 import warnings
 import pandas as pd
 
@@ -22,7 +15,6 @@ from urllib.error import URLError
 from unittest import mock
 from unittest.mock import mock_open
 
-from dirty_cat.utils import Version
 from dirty_cat.datasets import fetching
 from dirty_cat.datasets.utils import get_data_dir as _get_data_dir
 from dirty_cat.datasets.fetching import (
@@ -57,7 +49,8 @@ def test_fetch_openml_dataset():
     Though, to avoid the test being too long,
     we will download a small dataset (<1000 entries).
 
-    Reference: https://github.com/scikit-learn/scikit-learn/blob/master/sklearn/datasets/tests/test_openml.py
+    Reference:
+    https://github.com/scikit-learn/scikit-learn/blob/master/sklearn/datasets/tests/test_openml.py
     """
 
     # Hard-coded information about the test dataset.
@@ -98,7 +91,7 @@ def test_fetch_openml_dataset():
             # ``.gz`` files within in order to finish the test.
             pytest.skip(
                 'Exception: Skipping this test because we encountered an '
-                'issue probably due to an Internet connection problem.'
+                'issue probably related to an Internet connection problem. '
             )
             return
 
@@ -118,7 +111,7 @@ def test_fetch_openml_dataset():
         shutil.rmtree(path=str(test_data_dir), ignore_errors=True)
 
 
-@mock.patch("pathlib.Path.is_file")
+@mock.patch("dirty_cat.datasets.fetching.Path.is_file")
 @mock.patch("dirty_cat.datasets.fetching._get_features")
 @mock.patch("dirty_cat.datasets.fetching._get_details")
 @mock.patch("dirty_cat.datasets.fetching._export_gz_data_to_csv")
@@ -142,7 +135,7 @@ def test_fetch_openml_dataset_mocked(mock_download, mock_export,
 
     test_data_dir = get_test_data_dir()
 
-    # We test the function to see if it behaves correctly when files exists
+    # We test the function to see if it behaves correctly when files exist
     mock_pathlib_path_isfile.return_value = True
 
     fetch_openml_dataset(50, test_data_dir)
@@ -155,7 +148,7 @@ def test_fetch_openml_dataset_mocked(mock_download, mock_export,
     # Reset mocks
     mock_get_details.reset_mock()
 
-    # This time, the files does not exists
+    # This time, the files do not exist
     mock_pathlib_path_isfile.return_value = False
 
     fetch_openml_dataset(50, test_data_dir)
@@ -168,23 +161,21 @@ def test_fetch_openml_dataset_mocked(mock_download, mock_export,
     mock_get_details.assert_called_once()
 
 
-@mock.patch('sklearn.datasets.fetch_openml')
+@mock.patch('dirty_cat.datasets.fetching.fetch_openml')
 def test__download_and_write_openml_dataset(mock_fetch_openml):
     """Tests function ``_download_and_write_openml_dataset()``."""
 
+    dataset_id = 2
+
     test_data_dir = get_test_data_dir()
-    _download_and_write_openml_dataset(1, test_data_dir)
+    _download_and_write_openml_dataset(dataset_id, test_data_dir)
 
-    if Version(sklearn.__version__) >= Version('0.22'):
-        mock_fetch_openml.assert_called_once_with(data_id=1,
-                                                  data_home=str(test_data_dir),
-                                                  as_frame=True)
-    else:
-        mock_fetch_openml.assert_called_once_with(data_id=1,
-                                                  data_home=str(test_data_dir))
+    mock_fetch_openml.assert_called_once_with(data_id=dataset_id,
+                                              data_home=str(test_data_dir),
+                                              as_frame=True)
 
 
-@mock.patch("pathlib.Path.is_file")
+@mock.patch("dirty_cat.datasets.fetching.Path.is_file")
 def test__read_json_from_gz(mock_pathlib_path_isfile):
     """Tests function ``_read_json_from_gz()``."""
 
@@ -287,8 +278,10 @@ def test__export_gz_data_to_csv():
         with mock.patch("gzip.open",
                         mock_open(read_data=arff_data)) as mock_gzip_open:
             _export_gz_data_to_csv(dummy_gz, dummy_csv, features)
-            mock_pathlib_path_open.assert_called_with(mode='w', encoding='utf8')
-            mock_gzip_open.assert_called_with(dummy_gz, mode='rt', encoding='utf8')
+            mock_pathlib_path_open.assert_called_with(mode='w',
+                                                      encoding='utf8')
+            mock_gzip_open.assert_called_with(dummy_gz, mode='rt',
+                                              encoding='utf8')
 
 
 def test__features_to_csv_format():
@@ -338,8 +331,8 @@ def test_import_all_datasets(mock_fetch_dataset_as_dataclass,
     ):
         mock_fetch_dataset_as_dataclass.return_value = expected_return_value
 
-        returned_value = fetching.fetch_employee_salaries(drop_linked=False,
-                                                          drop_irrelevant=False)
+        returned_value = fetching.fetch_employee_salaries(
+            drop_linked=False, drop_irrelevant=False)
         assert expected_return_value == returned_value
 
         mock_fetch_openml_dataset.reset_mock()
diff --git a/dirty_cat/datasets/tests/test_utils.py b/dirty_cat/datasets/tests/test_utils.py
index 0b0bf84..811cd8c 100644
--- a/dirty_cat/datasets/tests/test_utils.py
+++ b/dirty_cat/datasets/tests/test_utils.py
@@ -1,13 +1,12 @@
-import os
-
 from pathlib import Path
-
 from unittest import mock
 
 
 @mock.patch('os.path.dirname')
 def test_get_data_dir(mock_os_path_dirname):
-    """Tests function ``get_data_dir()``."""
+    """
+    Tests function ``get_data_dir()``.
+    """
     from dirty_cat.datasets.utils import get_data_dir
 
     expected_return_value_default = Path("/user/directory/data")
diff --git a/dirty_cat/test/test_datetime_encoder.py b/dirty_cat/test/test_datetime_encoder.py
index 10e1630..0f4f02f 100644
--- a/dirty_cat/test/test_datetime_encoder.py
+++ b/dirty_cat/test/test_datetime_encoder.py
@@ -3,34 +3,35 @@ import pandas as pd
 from dirty_cat.datetime_encoder import DatetimeEncoder
 
 
-def get_date_array():
+def get_date_array() -> np.array:
     return np.array([pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03"]),
                      pd.to_datetime(["2021-02-03", "2020-02-04", "2021-02-05"]),
                      pd.to_datetime(["2022-01-01", "2020-12-25", "2022-01-03"]),
                      pd.to_datetime(["2023-02-03", "2020-02-04", "2023-02-05"])])
 
-def get_constant_date_array():
+
+def get_constant_date_array() -> np.array:
     return np.array([pd.to_datetime(["2020-01-01", "2020-02-04", "2021-02-05"]),
                      pd.to_datetime(["2020-01-01", "2020-02-04", "2021-02-05"]),
                      pd.to_datetime(["2020-01-01", "2020-02-04", "2021-02-05"]),
                      pd.to_datetime(["2020-01-01", "2020-02-04", "2021-02-05"])])
 
 
-def get_datetime_array():
+def get_datetime_array() -> np.array:
     return np.array([pd.to_datetime(["2020-01-01 10:12:01", "2020-01-02 10:23:00", "2020-01-03 10:00:00"]),
                      pd.to_datetime(["2021-02-03 12:45:23", "2020-02-04 22:12:00", "2021-02-05 12:00:00"]),
                      pd.to_datetime(["2022-01-01 23:23:43", "2020-12-25 11:12:00", "2022-01-03 11:00:00"]),
                      pd.to_datetime(["2023-02-03 11:12:12", "2020-02-04 08:32:00", "2023-02-05 23:00:00"])])
 
 
-def get_dirty_datetime_array():
+def get_dirty_datetime_array() -> np.array:
     return np.array([np.array(pd.to_datetime(["2020-01-01 10:12:01", "2020-01-02 10:23:00", "2020-01-03 10:00:00"])),
                      np.array(pd.to_datetime([np.nan, "2020-02-04 22:12:00", "2021-02-05 12:00:00"])),
                      np.array(pd.to_datetime(["2022-01-01 23:23:43", "2020-12-25 11:12:00", pd.NaT])),
                      np.array(pd.to_datetime(["2023-02-03 11:12:12", "2020-02-04 08:32:00", "2023-02-05 23:00:00"]))])
 
 
-def get_datetime_with_TZ_array():
+def get_datetime_with_TZ_array() -> pd.DataFrame:
     res = pd.DataFrame([pd.to_datetime(["2020-01-01 10:12:01"]),
                         pd.to_datetime(["2021-02-03 12:45:23"]),
                         pd.to_datetime(["2022-01-01 23:23:43"]),
@@ -40,7 +41,7 @@ def get_datetime_with_TZ_array():
     return res
 
 
-def test_fit():
+def test_fit() -> None:
     # Dates
     X = get_date_array()
     enc = DatetimeEncoder()
@@ -66,9 +67,11 @@ def test_fit():
     X = get_datetime_array()
     enc = DatetimeEncoder(add_day_of_the_week=True)
     expected_to_extract = ["year", "month", "day", "hour", "dayofweek"]
-    expected_features_per_column_ = {0: ["year", "month", "day", "hour", "dayofweek", "total_time"],
-                                     1: ["month", "day", "hour", "dayofweek", "total_time"],
-                                     2: ["year", "month", "day", "hour", "dayofweek"]}
+    expected_features_per_column_ = {
+        0: ["year", "month", "day", "hour", "dayofweek", "total_time"],
+        1: ["month", "day", "hour", "dayofweek", "total_time"],
+        2: ["year", "month", "day", "hour", "dayofweek"]
+    }
     enc.fit(X)
     assert enc._to_extract == expected_to_extract
     assert enc.features_per_column_ == expected_features_per_column_
@@ -76,9 +79,11 @@ def test_fit():
     X = get_datetime_array()
     enc = DatetimeEncoder(extract_until="minute")
     expected_to_extract = ["year", "month", "day", "hour", "minute"]
-    expected_features_per_column_ = {0: ["year", "month", "day", "hour", "minute", "total_time"],
-                                     1: ["month", "day", "hour", "minute"],
-                                     2: ["year", "month", "day", "hour"]}
+    expected_features_per_column_ = {
+        0: ["year", "month", "day", "hour", "minute", "total_time"],
+        1: ["month", "day", "hour", "minute"],
+        2: ["year", "month", "day", "hour"]
+    }
     enc.fit(X)
     assert enc._to_extract == expected_to_extract
     assert enc.features_per_column_ == expected_features_per_column_
@@ -87,9 +92,11 @@ def test_fit():
     X = get_dirty_datetime_array()
     enc = DatetimeEncoder()
     expected_to_extract = ["year", "month", "day", "hour"]
-    expected_features_per_column_ = {0: ["year", "month", "day", "hour", "total_time"],
-                                     1: ["month", "day", "hour", "total_time"],
-                                     2: ["year", "month", "day", "hour"]}
+    expected_features_per_column_ = {
+        0: ["year", "month", "day", "hour", "total_time"],
+        1: ["month", "day", "hour", "total_time"],
+        2: ["year", "month", "day", "hour"]
+    }
     enc.fit(X)
     assert enc._to_extract == expected_to_extract
     assert enc.features_per_column_ == expected_features_per_column_
@@ -98,7 +105,8 @@ def test_fit():
     X = get_datetime_with_TZ_array()
     enc = DatetimeEncoder()
     expected_to_extract = ["year", "month", "day", "hour"]
-    expected_features_per_column_ = {0: ["year", "month", "day", "hour", "total_time"]}
+    expected_features_per_column_ = {0: ["year", "month", "day",
+                                         "hour", "total_time"]}
     enc.fit(X)
     assert enc._to_extract == expected_to_extract
     assert enc.features_per_column_ == expected_features_per_column_
@@ -107,9 +115,11 @@ def test_fit():
     # Without column names
     X = get_datetime_array()
     enc = DatetimeEncoder(add_day_of_the_week=True)
-    expected_feature_names = ["0_year", "0_month", "0_day", "0_hour", "0_dayofweek", "0_total_time",
-                              "1_month", "1_day", "1_hour", "1_dayofweek", "1_total_time",
-                              "2_year", "2_month", "2_day", "2_hour", "2_dayofweek"]
+    expected_feature_names = [
+        "0_year", "0_month", "0_day", "0_hour", "0_dayofweek", "0_total_time",
+        "1_month", "1_day", "1_hour", "1_dayofweek", "1_total_time",
+        "2_year", "2_month", "2_day", "2_hour", "2_dayofweek"
+    ]
     enc.fit(X)
     assert enc.get_feature_names_out() == expected_feature_names
 
@@ -118,9 +128,11 @@ def test_fit():
     X = pd.DataFrame(X)
     X.columns = ["col1", "col2", "col3"]
     enc = DatetimeEncoder(add_day_of_the_week=True)
-    expected_feature_names = ["col1_year", "col1_month", "col1_day", "col1_hour", "col1_dayofweek", "col1_total_time",
-                              "col2_month", "col2_day", "col2_hour", "col2_dayofweek", "col2_total_time",
-                              "col3_year", "col3_month", "col3_day", "col3_hour", "col3_dayofweek"]
+    expected_feature_names = [
+        "col1_year", "col1_month", "col1_day", "col1_hour", "col1_dayofweek", "col1_total_time",
+        "col2_month", "col2_day", "col2_hour", "col2_dayofweek", "col2_total_time",
+        "col3_year", "col3_month", "col3_day", "col3_hour", "col3_dayofweek"
+    ]
     enc.fit(X)
     assert enc.get_feature_names_out() == expected_feature_names
 
@@ -160,7 +172,8 @@ def test_transform():
                                 [2021, 2, 3, 12, 2, 0],
                                 [2022, 1, 1, 23, 5, 0],
                                 [2023, 2, 3, 11, 4, 0]]).astype(np.float64)
-    expected_result[:, 5] = (X.astype('int64')//1e9).astype(np.float64).reshape(-1) #time from epochs in seconds
+    # Time from epochs in seconds
+    expected_result[:, 5] = (X.astype('int64')//1e9).astype(np.float64).reshape(-1)
 
     enc.fit(X)
     X_trans = enc.transform(X)
@@ -169,39 +182,47 @@ def test_transform():
     # Check if we find back the date from the time to epoch
     assert ((pd.to_datetime(X_trans[:, 5], unit="s") - pd.to_datetime(X.reshape(-1))).total_seconds() == 0).all()
 
-    # Dirty Datetimes
+    # Dirty datetimes
     X = get_dirty_datetime_array()[:, 0].reshape(-1, 1)
     enc = DatetimeEncoder(add_day_of_the_week=True)
     expected_result = np.array([[2020, 1, 1, 10, 2, 0],
-                                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
+                                [np.nan] * 6,
                                 [2022, 1, 1, 23, 5, 0],
                                 [2023, 2, 3, 11, 4, 0]])
-    expected_result[:, 5] = (X.astype('int64')//1e9).astype(np.float64).reshape(-1) #time from epochs in seconds
+    # Time from epochs in seconds
+    expected_result[:, 5] = (X.astype('int64')//1e9).astype(np.float64).reshape(-1)
     expected_result[1, 5] = np.nan
     enc.fit(X)
     X_trans = enc.transform(X)
     assert np.allclose(X_trans, expected_result, equal_nan=True)
 
     # Datetimes with TZ
-    # If the dates are timezone-aware, all the feature extractions should be done
-    # in the provided timezone.
-    # But the full time to epoch should correspond to the true number of seconds between epoch time and the time
-    # of the date.
+    # If the dates are timezone-aware, all the feature extractions should
+    # be done in the provided timezone.
+    # But the full time to epoch should correspond to the true number of
+    # seconds between epoch time and the time of the date.
     X = get_datetime_with_TZ_array()
     enc = DatetimeEncoder(add_day_of_the_week=True)
     expected_result = np.array([[2020, 1, 1, 10, 2, 0],
                                 [2021, 2, 3, 12, 2, 0],
                                 [2022, 1, 1, 23, 5, 0],
                                 [2023, 2, 3, 11, 4, 0]]).astype(np.float64)
-    expected_result[:, 5] = (X.iloc[:, 0].view(dtype='int64')//1e9).astype(np.float64).to_numpy().reshape(-1) #time from epochs in seconds
+    # Time from epochs in seconds
+    expected_result[:, 5] = (X.iloc[:, 0].view(dtype='int64')//1e9).astype(np.float64).to_numpy().reshape(-1)
     enc.fit(X)
     X_trans = enc.transform(X)
     assert np.allclose(X_trans, expected_result, equal_nan=True)
 
     # Check if we find back the date from the time to epoch
-    assert ((pd.to_datetime(X_trans[:, 5], unit="s").tz_localize("utc").tz_convert(X.iloc[:, 0][0].tz) - pd.DatetimeIndex(X.iloc[:, 0])).total_seconds() == 0).all()
+    assert (
+            (
+                pd.to_datetime(X_trans[:, 5], unit="s")
+                .tz_localize("utc")
+                .tz_convert(X.iloc[:, 0][0].tz) - pd.DatetimeIndex(X.iloc[:, 0])
+            ).total_seconds() == 0
+    ).all()
 
     # Check if it's working when the date is constant
     X = get_constant_date_array()
     enc = DatetimeEncoder(add_day_of_the_week=True)
-    assert enc.fit_transform(X).shape[1] == 0
\ No newline at end of file
+    assert enc.fit_transform(X).shape[1] == 0
diff --git a/dirty_cat/test/test_fast_hash.py b/dirty_cat/test/test_fast_hash.py
index 01cf0a0..ed478b2 100644
--- a/dirty_cat/test/test_fast_hash.py
+++ b/dirty_cat/test/test_fast_hash.py
@@ -1,10 +1,10 @@
 from dirty_cat.fast_hash import ngram_min_hash
+from sklearn.datasets import fetch_20newsgroups
 
 
 def test_fast_hash():
 
-    from sklearn import datasets
-    data = datasets.fetch_20newsgroups()
+    data = fetch_20newsgroups()
     a = data.data[0]
 
     min_hash = ngram_min_hash(a, seed=0)
diff --git a/dirty_cat/test/test_gap_encoder.py b/dirty_cat/test/test_gap_encoder.py
index a77f983..4e5aa80 100644
--- a/dirty_cat/test/test_gap_encoder.py
+++ b/dirty_cat/test/test_gap_encoder.py
@@ -1,21 +1,27 @@
 import time
-import numpy as np
 import pytest
+import numpy as np
 import pandas as pd
 
+from typing import Literal
+
 from sklearn import __version__ as sklearn_version
 from sklearn.datasets import fetch_20newsgroups
+
+from dirty_cat.datasets import fetch_employee_salaries
 from dirty_cat.utils import Version
 from dirty_cat import GapEncoder
 
-@pytest.mark.parametrize("init1, analyzer1, analyzer2",[
+
+@pytest.mark.parametrize("init1, analyzer1, analyzer2", [
     ('k-means++', 'char', 'word'),
     ('random', 'char', 'word'),
     ('k-means', 'char', 'word')
 ])
-def test_analyzer(init1, analyzer1, analyzer2):
-    """" Test if the output is different when the analyzer is 'word' or 'char'.
-        If it is, no error ir raised. 
+def test_analyzer(init1: str, analyzer1: str, analyzer2: str) -> None:
+    """
+    Test if the output is different when the analyzer is 'word' or 'char'.
+    If it is, no error ir raised.
     """
     add_words = False
     n_samples = 70
@@ -38,15 +44,18 @@ def test_analyzer(init1, analyzer1, analyzer2):
     encoder.fit(X)
     y2 = encoder.transform(X)
     
-    # Test inequality btw analyzer word and char ouput:
-    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, y, y2)
+    # Test inequality btw analyzer word and char output:
+    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal,
+                             y, y2)
+
 
 @pytest.mark.parametrize("hashing, init, analyzer, add_words", [
     (False, 'k-means++', 'word', True),
     (True, 'random', 'char', False),
     (True, 'k-means', 'char_wb', True)
 ])
-def test_gap_encoder(hashing, init, analyzer, add_words, n_samples=70):
+def test_gap_encoder(hashing: bool, init: str, analyzer: str, add_words: bool,
+                     n_samples: int = 70) -> None:
     X_txt = fetch_20newsgroups(subset='train')['data'][:n_samples]
     X = np.array([X_txt, X_txt]).T
     n_components = 10
@@ -75,7 +84,7 @@ def test_gap_encoder(hashing, init, analyzer, add_words, n_samples=70):
     np.testing.assert_array_equal(y, y2)
 
 
-def test_input_type():
+def test_input_type() -> None:
     # Numpy array with one column
     X = np.array([['alice'], ['bob']])
     enc = GapEncoder(n_components=2, random_state=42)
@@ -97,10 +106,9 @@ def test_input_type():
     X_enc_df = enc.fit_transform(df)
     # Check if the encoded vectors are the same
     np.testing.assert_array_equal(X_enc_array, X_enc_df)
-    return
 
 
-def test_partial_fit(n_samples=70):
+def test_partial_fit(n_samples=70) -> None:
     X_txt = fetch_20newsgroups(subset='train')['data'][:n_samples]
     X = np.array([X_txt, X_txt]).T
     # Gap encoder with fit on one batch
@@ -112,10 +120,9 @@ def test_partial_fit(n_samples=70):
     X_enc_partial = enc.transform(X)
     # Check if the encoded vectors are the same
     np.testing.assert_almost_equal(X_enc, X_enc_partial)
-    return
 
 
-def test_get_feature_names_out(n_samples=70):
+def test_get_feature_names_out(n_samples=70) -> None:
     X_txt = fetch_20newsgroups(subset='train')['data'][:n_samples]
     X = np.array([X_txt, X_txt]).T
     enc = GapEncoder(random_state=42)
@@ -138,17 +145,16 @@ def test_get_feature_names_out(n_samples=70):
     return
 
 
-def test_overflow_error():
+def test_overflow_error() -> None:
     np.seterr(over='raise', divide='raise')
     r = np.random.RandomState(0)
     X = r.randint(1e5, 1e6, size=(8000, 1)).astype(str)
     enc = GapEncoder(n_components=2, batch_size=1, min_iter=1, max_iter=1,
                      random_state=0)
     enc.fit(X)
-    return
 
 
-def test_score(n_samples=70):
+def test_score(n_samples: int = 70) -> None:
     X_txt = fetch_20newsgroups(subset='train')['data'][:n_samples]
     X1 = np.array(X_txt)[:, None]
     X2 = np.hstack([X1, X1])
@@ -159,36 +165,35 @@ def test_score(n_samples=70):
     score_X2 = enc.score(X2)
     # Check that two identical columns give the same score
     assert score_X1 * 2 == score_X2
-    return
 
 
 @pytest.mark.parametrize("missing", ['zero_impute', 'error', 'aaa'])
-def test_missing_values(missing):
+def test_missing_values(missing: str) -> None:
     observations = [['alice', 'bob'], ['bob', 'alice'], ['bob', np.nan],
                     ['alice', 'charlie'], [np.nan, 'alice']]
     observations = np.array(observations, dtype=object)
     enc = GapEncoder(handle_missing=missing, n_components=3)
     if missing == 'error':
-        with pytest.raises(ValueError, match=r'Input data contains missing values.'):
+        with pytest.raises(ValueError,
+                           match='Input data contains missing values'):
             enc.fit_transform(observations)
     elif missing == 'zero_impute':
         enc.fit_transform(observations)
         enc.partial_fit(observations)
     else:
-        with pytest.raises(ValueError, match=r"handle_missing should be either "
-                                             r"'error' or 'zero_impute', got 'aaa'"):
+        with pytest.raises(ValueError,
+                           match=r"handle_missing should be either "
+                                 r"'error' or 'zero_impute', got 'aaa'"):
             enc.fit_transform(observations)
 
 
-def profile_encoder(Encoder, init):
-    # not an unit test
-
-    from dirty_cat.datasets import fetch_employee_salaries
+def profile_encoder(init: str) -> None:
+    # not a unit test
     info = fetch_employee_salaries()
     data = pd.read_csv(info['path'], **info['read_csv_kwargs'])
     X = np.array(data['employee_position_title'])[:, None]
     t0 = time.time()
-    encoder = Encoder(n_components=50, init=init)
+    encoder = GapEncoder(n_components=50, init=init)
     encoder.fit(X)
     y = encoder.transform(X)
     assert y.shape == (len(X), 50)
@@ -200,33 +205,38 @@ if __name__ == '__main__':
     
     print('test_analyzer')
     test_analyzer('k-means++', 'char_wb', 'word')
-    print('test_analyzer passed')
+    print('passed')
+
     print('start test_gap_encoder')
-    test_gap_encoder(True, 'k-means++', 'char', False)
-    print('test_gap_encoder passed')
+    test_gap_encoder(hashing=True, init='k-means++', analyzer='char', add_words=False)
+    print('passed')
+
     print('start test_input_type')
     test_input_type()
-    print('test_input_type passed')
+    print('passed')
+
     print('start test_partial_fit')
     test_partial_fit()
-    print('test_partial_fit passed')
+    print('passed')
+
     print('start test_get_feature_names_out')
     test_get_feature_names_out()
-    print('test_get_feature_names_out passed')
+    print('passed')
+
     print('start test_overflow_error')
     test_overflow_error()
-    print('test_overflow_error passed')
+    print('passed')
+
     print('start test_score')
     test_score()
     print('test_score passed')
     
     for _ in range(3):
         print('time profile_encoder(GapEncoder, init="k-means++")')
-        print("{:.4} seconds".format(profile_encoder(
-            GapEncoder, init='k-means++')))
+        print(f"{profile_encoder('k-means++'):.4} seconds")
+
     for _ in range(3):
         print('time profile_encoder(GapEncoder, init="k-means")')
-        print("{:.4} seconds".format(profile_encoder(
-            GapEncoder, init='k-means')))
+        print(f"{profile_encoder('k-means'):.4} seconds")
     
     print('Done')
diff --git a/dirty_cat/test/test_minhash_encoder.py b/dirty_cat/test/test_minhash_encoder.py
index 6283e3a..fc8745e 100644
--- a/dirty_cat/test/test_minhash_encoder.py
+++ b/dirty_cat/test/test_minhash_encoder.py
@@ -1,22 +1,25 @@
 import time
+import pytest
 import random
-from string import ascii_lowercase
 import numpy as np
 import pandas as pd
-import pytest
 
+from string import ascii_lowercase
 from sklearn.datasets import fetch_20newsgroups
+
 from dirty_cat import MinHashEncoder
+from dirty_cat.datasets import fetch_employee_salaries
+
 
-def test_MinHashEncoder(n_sample=70, minmax_hash=False):
+def test_MinHashEncoder(n_sample: int = 70) -> None:
     X_txt = fetch_20newsgroups(subset='train')['data']
-    X = np.array(X_txt[:n_sample])[:,None]
+    X = np.array(X_txt[:n_sample])[:, None]
 
     for minmax_hash in [True, False]:
         for hashing in ['fast', 'murmur']:
 
             if minmax_hash and hashing == 'murmur':
-                pass # not implemented
+                pass  # not implemented
 
             # Test output shape
             encoder = MinHashEncoder(n_components=50, hashing=hashing)
@@ -40,23 +43,28 @@ def test_MinHashEncoder(n_sample=70, minmax_hash=False):
                 y_substring = encoder.transform(X_substring)
                 np.testing.assert_array_less(y - y_substring, 0.0001)
 
-def test_multiple_columns():
-    """ This test is intented to verify that fitting multiple columns
-        with the MinHashEncoder will not produce an error, and will 
-        encode the column independently """
+
+def test_multiple_columns() -> None:
+    """
+    This test aims at verifying that fitting multiple columns
+    with the MinHashEncoder will not produce an error, and will
+    encode the column independently.
+    """
     X = pd.DataFrame([('bird', 'parrot'),
                       ('bird', 'nightingale'),
                       ('mammal', 'monkey'),
                       ('mammal', np.nan)],
-                      columns=('class', 'type'))
+                     columns=('class', 'type'))
     X1 = X[['class']]
     X2 = X[['type']]
     fit1 = MinHashEncoder(n_components=30).fit_transform(X1)
     fit2 = MinHashEncoder(n_components=30).fit_transform(X2)
     fit = MinHashEncoder(n_components=30).fit_transform(X)
-    assert np.array_equal(np.array([fit[:, :30], fit[:, 30:60]]), np.array([fit1, fit2]))
+    assert np.array_equal(np.array([fit[:, :30], fit[:, 30:60]]),
+                          np.array([fit1, fit2]))
 
-def test_input_type():
+
+def test_input_type() -> None:
     # Numpy array
     X = np.array(['alice', 'bob'])[:,None]
     enc = MinHashEncoder(n_components=2)
@@ -66,28 +74,15 @@ def test_input_type():
     enc = MinHashEncoder(n_components=2)
     enc.fit_transform(X)
 
-def test_get_unique_ngrams():
-    string = 'test'
-    true_ngrams = {
-        (' ','t'), ('t','e'), ('e','s'), ('s', 't'),
-        ('t',' '), (' ','t','e'), ('t','e','s'),
-        ('e','s','t'), ('s','t',' '), (' ','t','e','s'),
-        ('t','e','s','t'), ('e','s','t',' ')}
-    ngram_range = (2,4)
-    enc = MinHashEncoder(n_components=2)
-    ngrams = enc.get_unique_ngrams(string, ngram_range)
-    assert ngrams == true_ngrams
-
 
-def profile_encoder(encoder, hashing='fast', minmax_hash=False):
-    # not an unit test
-
-    from dirty_cat.datasets import fetch_employee_salaries
+def profile_encoder(hashing: str = 'fast', minmax_hash: bool = False) -> float:
+    # not a unit test
     employee_salaries = fetch_employee_salaries()
     df = employee_salaries.X
     X = df[["employee_position_title"]]
     t0 = time.time()
-    enc = encoder(n_components=50, hashing=hashing, minmax_hash=minmax_hash)
+    enc = MinHashEncoder(n_components=50, hashing=hashing,
+                         minmax_hash=minmax_hash)
     enc.fit(X)
     y = enc.transform(X)
     assert y.shape == (len(X), 50)
@@ -99,7 +94,7 @@ def profile_encoder(encoder, hashing='fast', minmax_hash=False):
     ['numpy', 'error', 'fast'],
     ['pandas', 'zero_impute', 'murmur'],
     ['numpy', 'zero_impute', 'fast']])
-def test_missing_values(input_type, missing, hashing):
+def test_missing_values(input_type: str, missing: str, hashing: str) -> None:
     X = ['Red',
          np.nan,
          'green',
@@ -114,7 +109,6 @@ def test_missing_values(input_type, missing, hashing):
     if input_type == 'numpy':
         X = np.array(X, dtype=object)[:,None]
     elif input_type == 'pandas':
-        pd = pytest.importorskip("pandas")
         X = pd.DataFrame(X)
 
     encoder = MinHashEncoder(n_components=n, hashing=hashing,
@@ -122,8 +116,7 @@ def test_missing_values(input_type, missing, hashing):
     if missing == 'error':
         encoder.fit(X)
         if input_type in ['numpy', 'pandas']:
-            with pytest.raises(ValueError, match=r"missing"
-                               " values in input"):
+            with pytest.raises(ValueError, match=r"missing values in input"):
                 encoder.transform(X)
     elif missing == 'zero_impute':
         encoder.fit(X)
@@ -134,47 +127,45 @@ def test_missing_values(input_type, missing, hashing):
             assert np.array_equal(y[1], z)
             assert np.array_equal(y[-1], z)
     else:
-        with pytest.raises(ValueError, match=r"handle_missing"
-                           " should be either 'error' or 'zero_impute'"):
+        with pytest.raises(ValueError, match=r"expected any of"):
             encoder.fit_transform(X)
     return
 
 
-def test_cache_overflow():
+def test_cache_overflow() -> None:
     # Regression test for cache overflow resulting in -1s in encoding
     def get_random_string(length):
         letters = ascii_lowercase
-        result_str = ''.join(random.choice(letters) for i in range(length))
+        result_str = ''.join(random.choice(letters) for _ in range(length))
         return result_str
 
     encoder = MinHashEncoder(n_components=3)
     capacity = encoder._capacity
-    raw_data = [get_random_string(10) for x in range(capacity + 1)]
-    raw_data = np.array(raw_data)[:,None]
+    raw_data = [get_random_string(10) for _ in range(capacity + 1)]
+    raw_data = np.array(raw_data)[:, None]
     y = encoder.fit_transform(raw_data)
 
     assert len(y[y == -1.0]) == 0
 
 
 if __name__ == '__main__':
-    print('start test')
+    print('start test_MinHashEncoder')
     test_MinHashEncoder()
-    print('test passed')
+    print('passed')
     
-    print('start test')
+    print('start test_multiple_columns')
     test_multiple_columns()
-    print('multiple columns encoding test passed')
+    print('passed')
 
     for _ in range(3):
         print('time profile_encoder(MinHashEncoder, hashing=fast)')
-        print("{:.4} seconds".format(
-            profile_encoder(MinHashEncoder, hashing='fast')))
+        print(f"{profile_encoder(hashing='fast'):.4} seconds")
+
     for _ in range(3):
         print('time profile_encoder(MinHashEncoder, hashing=fast) with minmax')
-        print("{:.4} seconds".format(profile_encoder(MinHashEncoder,
-                                     hashing='fast', minmax_hash=True)))
+        print(f"{profile_encoder(hashing='fast', minmax_hash=True):.4} seconds")
+
     print('time profile_encoder(MinHashEncoder, hashing=murmur)')
-    print("{:.4} seconds".format(
-        profile_encoder(MinHashEncoder, hashing='murmur')))
+    print(f"{profile_encoder(hashing='murmur'):.4} seconds")
 
     print('Done')
diff --git a/dirty_cat/test/test_similarity_encoder.py b/dirty_cat/test/test_similarity_encoder.py
index 2c547b9..2cbe52c 100644
--- a/dirty_cat/test/test_similarity_encoder.py
+++ b/dirty_cat/test/test_similarity_encoder.py
@@ -2,6 +2,8 @@ import numpy as np
 import numpy.testing
 import pytest
 
+from typing import Callable, Optional
+
 from sklearn import __version__ as sklearn_version
 
 from dirty_cat.utils import Version
@@ -9,7 +11,7 @@ from dirty_cat import similarity_encoder, string_distances
 from dirty_cat.similarity_encoder import get_kmeans_prototypes
 
 
-def test_specifying_categories():
+def test_specifying_categories() -> None:
     # When creating a new SimilarityEncoder:
     # - if categories = 'auto', the categories are the sorted, unique training
     # set observations (for each column)
@@ -19,7 +21,7 @@ def test_specifying_categories():
     # In this test, we first find the sorted, unique categories in the training
     # set, and create a SimilarityEncoder by giving it explicitly the computed
     # categories. The test consists in making sure the transformed observations
-    # given by this encoder are equal to the transformed obervations in the
+    # given by this encoder are equal to the transformed observations in the
     # case of a SimilarityEncoder created with categories = 'auto'
 
     observations = [['bar'], ['foo']]
@@ -35,7 +37,7 @@ def test_specifying_categories():
     assert np.allclose(feature_matrix_auto_cat, feature_matrix_with_cat)
 
 
-def test_fast_ngram_similarity():
+def test_fast_ngram_similarity() -> None:
     vocabulary = [['bar', 'foo']]
     observations = [['foo'], ['baz']]
 
@@ -66,21 +68,18 @@ def _test_missing_values(input_type, missing):
 
     sim_enc = similarity_encoder.SimilarityEncoder(handle_missing=missing)
     if missing == 'error':
-        with pytest.raises(ValueError, match=r"Found missing values in input "
-                           "data; set handle_missing='' to encode "
-                           "with missing values"):
+        with pytest.raises(ValueError, match=r"Found missing values in input"):
             sim_enc.fit_transform(observations)
     elif missing == '':
         ans = sim_enc.fit_transform(observations)
         assert np.allclose(encoded, ans)
     else:
-        with pytest.raises(ValueError, match=r"handle_missing"
-                           " should be either 'error' or ''"):
+        with pytest.raises(ValueError, match=r"expected any of"):
             sim_enc.fit_transform(observations)
         return
 
 
-def _test_missing_values_transform(input_type, missing):
+def _test_missing_values_transform(input_type: str, missing: str) -> None:
     observations = [['a', 'b'], ['b', 'a'], ['b', 'b'],
                     ['a', 'c'], ['c', 'a']]
     test_observations = [['a', 'b'], ['b', 'a'], ['b', np.nan],
@@ -100,9 +99,7 @@ def _test_missing_values_transform(input_type, missing):
     sim_enc = similarity_encoder.SimilarityEncoder(handle_missing=missing)
     if missing == 'error':
         sim_enc.fit_transform(observations)
-        with pytest.raises(ValueError, match=r"Found missing values in input "
-                           "data; set handle_missing='' to encode "
-                           "with missing values"):
+        with pytest.raises(ValueError, match=r"Found missing values in input"):
             sim_enc.transform(test_observations)
     elif missing == '':
         sim_enc.fit_transform(observations)
@@ -110,7 +107,9 @@ def _test_missing_values_transform(input_type, missing):
         assert np.allclose(encoded, ans)
 
 
-def _test_similarity(similarity_f, hashing_dim=None, categories='auto', n_prototypes=None):
+def _test_similarity(similarity_f: Callable, hashing_dim: Optional[int] = None,
+                     categories: str = 'auto', n_prototypes: int = None,
+                     ) -> None:
     if n_prototypes is None:
         X = np.array(['aa', 'aaa', 'aaab']).reshape(-1, 1)
         X_test = np.array([['Aa', 'aAa', 'aaa', 'aaab', ' aaa  c']]).reshape(-1, 1)
@@ -160,7 +159,7 @@ def _test_similarity(similarity_f, hashing_dim=None, categories='auto', n_protot
         numpy.testing.assert_almost_equal(encoder, ans)
 
 
-def test_similarity_encoder():
+def test_similarity_encoder() -> None:
     categories = ['auto', 'most_frequent', 'k-means']
     for category in categories:
         if category == 'auto':
@@ -184,21 +183,23 @@ def test_similarity_encoder():
             _test_missing_values_transform(input_type, missing)
 
 
-def test_kmeans_protoypes():
+def test_kmeans_protoypes() -> None:
     X_test = np.array(['cbbba', 'baaac', 'accc'])
     proto = get_kmeans_prototypes(X_test, 3)
     assert np.array_equal(np.sort(proto), np.sort(X_test))
 
 
-def test_reproducibility():
-    sim_enc = similarity_encoder.SimilarityEncoder(categories='k-means', n_prototypes=10, random_state=435)
+def test_reproducibility() -> None:
+    sim_enc = similarity_encoder.SimilarityEncoder(
+        categories='k-means', n_prototypes=10, random_state=435,
+    )
     X = np.array([' %s ' % chr(i) for i in range(32, 127)]).reshape((-1, 1))
     prototypes = sim_enc.fit(X).categories_[0]
     for i in range(10):
         assert (np.array_equal(prototypes, sim_enc.fit(X).categories_[0]))
 
 
-def test_get_features():
+def test_get_features() -> None:
     # See https://github.com/dirty-cat/dirty_cat/issues/168
     sim_enc = similarity_encoder.SimilarityEncoder(random_state=435)
     X = np.array(['%s' % chr(i) for i in range(32, 127)]).reshape((-1, 1))
@@ -220,4 +221,3 @@ def test_get_features():
         'x0_q', 'x0_r', 'x0_s', 'x0_t', 'x0_u', 'x0_v', 'x0_w', 'x0_x', 'x0_y',
         'x0_z', 'x0_{', 'x0_|', 'x0_}', 'x0_~'
     ]
-
diff --git a/dirty_cat/test/test_sklearn.py b/dirty_cat/test/test_sklearn.py
index 5205739..0b626f8 100644
--- a/dirty_cat/test/test_sklearn.py
+++ b/dirty_cat/test/test_sklearn.py
@@ -1,6 +1,7 @@
-from sklearn.utils.estimator_checks import check_estimator, parametrize_with_checks
+from sklearn.utils.estimator_checks import check_estimator
 
-from dirty_cat import DatetimeEncoder, GapEncoder, MinHashEncoder, SimilarityEncoder, SuperVectorizer, TargetEncoder
+from dirty_cat import (DatetimeEncoder, GapEncoder, MinHashEncoder,
+                       SimilarityEncoder, SuperVectorizer, TargetEncoder)
 
 
 def test_sklearn_compatible_DatetimeEncoder():
@@ -24,4 +25,4 @@ def test_sklearn_compatible_SimilarityEncoder():
 
 
 def test_sklearn_compatible_TargetEncoder():
-    check_estimator(TargetEncoder())
\ No newline at end of file
+    check_estimator(TargetEncoder())
diff --git a/dirty_cat/test/test_string_distances.py b/dirty_cat/test/test_string_distances.py
index bbd66b9..3933369 100644
--- a/dirty_cat/test/test_string_distances.py
+++ b/dirty_cat/test/test_string_distances.py
@@ -1,9 +1,23 @@
 import numpy as np
 
+from typing import List, Tuple
 from dirty_cat import string_distances
 
 
-def _random_string_pairs(n_pairs=50, seed=1):
+def test_get_unique_ngrams() -> None:
+    string = 'test'
+    true_ngrams = {
+        (' ', 't'), ('t', 'e'), ('e', 's'), ('s', 't'),
+        ('t', ' '), (' ', 't', 'e'), ('t', 'e', 's'),
+        ('e', 's', 't'), ('s', 't', ' '), (' ', 't', 'e', 's'),
+        ('t', 'e', 's', 't'), ('e', 's', 't', ' ')
+    }
+    ngram_range = (2, 4)
+    ngrams = string_distances.get_unique_ngrams(string, ngram_range)
+    assert ngrams == true_ngrams
+
+
+def _random_string_pairs(n_pairs=50, seed=1) -> List[Tuple[str, str]]:
     rng = np.random.RandomState(seed)
     characters = list(map(chr, range(10000)))
     pairs = []
@@ -16,13 +30,13 @@ def _random_string_pairs(n_pairs=50, seed=1):
     return pairs
 
 
-def _check_symmetry(dist_func, *args, **kwargs):
+def _check_symmetry(dist_func, *args, **kwargs) -> None:
     for (a, b) in _random_string_pairs():
         assert dist_func(
             a, b, *args, **kwargs) == dist_func(b, a, *args, **kwargs)
 
 
-def test_ngram_similarity():
+def test_ngram_similarity() -> None:
     # TODO
     # assert ...
     for n in range(1, 4):
diff --git a/dirty_cat/test/test_super_vectorizer.py b/dirty_cat/test/test_super_vectorizer.py
index ad7b07d..0eb32dc 100644
--- a/dirty_cat/test/test_super_vectorizer.py
+++ b/dirty_cat/test/test_super_vectorizer.py
@@ -3,34 +3,39 @@ import pytest
 import sklearn
 import pandas as pd
 
+from typing import List, Dict
+
 from sklearn.preprocessing import StandardScaler
 from sklearn.utils.validation import check_is_fitted
 from sklearn.exceptions import NotFittedError
 
-
 from dirty_cat import SuperVectorizer
 from dirty_cat import GapEncoder
 from dirty_cat.utils import Version
 
 
-def check_same_transformers(expected_transformers: dict, actual_transformers: list):
+def check_same_transformers(expected_transformers: dict,
+                            actual_transformers: list):
     # Construct the dict from the actual transformers
-    actual_transformers_dict = dict([(name, cols) for name, trans, cols in actual_transformers])
+    actual_transformers_dict = {name: cols
+                                for name, trans, cols in actual_transformers}
     assert actual_transformers_dict == expected_transformers
 
+
 def type_equality(expected_type, actual_type):
     """
-    Checks that the expected type is equal to the actual type, assuming object and str types
-    are equivalent (considered as categorical by the supervectorizer).
+    Checks that the expected type is equal to the actual type,
+    assuming object and str types are equivalent
+    (considered as categorical by the SuperVectorizer).
     """
-    if (isinstance(expected_type , object) or  isinstance(expected_type, str))\
+    if (isinstance(expected_type, object) or isinstance(expected_type, str))\
             and (isinstance(actual_type, object) or isinstance(actual_type, str)):
         return True
     else:
         return expected_type == actual_type
 
 
-def _get_clean_dataframe():
+def _get_clean_dataframe() -> pd.DataFrame:
     """
     Creates a simple DataFrame with various types of data,
     and without missing values.
@@ -45,11 +50,11 @@ def _get_clean_dataframe():
     })
 
 
-def _get_dirty_dataframe():
+def _get_dirty_dataframe() -> pd.DataFrame:
     """
     Creates a simple DataFrame with some missing values.
     We'll use different types of missing values (np.nan, pd.NA, None)
-    to see how robust the vectorizer is.
+    to test the robustness of the vectorizer.
     """
     return pd.DataFrame({
         'int': pd.Series([15, 56, pd.NA, 12, 44], dtype='Int64'),
@@ -60,7 +65,8 @@ def _get_dirty_dataframe():
         'cat2': pd.Series(['20K+', '40K+', '60K+', '30K+', np.nan], dtype='object'),
     })
 
-def _get_numpy_array():
+
+def _get_numpy_array() -> np.array:
     return np.array([["15", "56", pd.NA, "12", ""],
             ["?", "2.4", "6.2", "10.45", np.nan],
             ['public', np.nan, 'private', 'private', pd.NA],
@@ -68,12 +74,15 @@ def _get_numpy_array():
             [np.nan, 'yes', 'no', 'yes', 'no'],
             ['20K+', '40K+', '60K+', '30K+', np.nan]]).T
 
-def _get_list_of_lists():
+
+def _get_list_of_lists() -> list:
     return _get_numpy_array().tolist()
 
-def _get_datetimes_dataframe():
+
+def _get_datetimes_dataframe() -> pd.DataFrame:
     """
-    Creates a DataFrame with various date format, already converted or to be converted.
+    Creates a DataFrame with various date formats,
+    already converted or to be converted.
     """
     return pd.DataFrame({
         "pd_datetime": [
@@ -112,12 +121,12 @@ def _get_datetimes_dataframe():
 
 def _test_possibilities(
         X,
-        expected_transformers_df,
-        expected_transformers_2,
-        expected_transformers_np_no_cast,
-        expected_transformers_series,
-        expected_transformers_plain,
-        expected_transformers_np_cast,
+        expected_transformers_df: Dict[str, List[str]],
+        expected_transformers_2: Dict[str, List[str]],
+        expected_transformers_np_no_cast: Dict[str, List[int]],
+        expected_transformers_series: Dict[str, List[str]],
+        expected_transformers_plain: Dict[str, List[str]],
+        expected_transformers_np_cast: Dict[str, List[int]],
 ):
     """
     Do a bunch of tests with the SuperVectorizer.
@@ -133,22 +142,26 @@ def _test_possibilities(
     )
     # Warning: order-dependant
     vectorizer_base.fit_transform(X)
-    check_same_transformers(expected_transformers_df, vectorizer_base.transformers)
+    check_same_transformers(expected_transformers_df,
+                            vectorizer_base.transformers)
 
     # Test with higher cardinality threshold and no numeric transformer
     vectorizer_default = SuperVectorizer()  # Using default values
     vectorizer_default.fit_transform(X)
-    check_same_transformers(expected_transformers_2, vectorizer_default.transformers)
+    check_same_transformers(expected_transformers_2,
+                            vectorizer_default.transformers)
 
     # Test with a numpy array
     arr = X.to_numpy()
     # Instead of the columns names, we'll have the column indices.
     vectorizer_base.fit_transform(arr)
-    check_same_transformers(expected_transformers_np_no_cast, vectorizer_base.transformers)
+    check_same_transformers(expected_transformers_np_no_cast,
+                            vectorizer_base.transformers)
 
     # Test with pandas series
     vectorizer_base.fit_transform(X['cat1'])
-    check_same_transformers(expected_transformers_series, vectorizer_base.transformers)
+    check_same_transformers(expected_transformers_series,
+                            vectorizer_base.transformers)
 
     # Test casting values
     vectorizer_cast = SuperVectorizer(
@@ -160,10 +173,12 @@ def _test_possibilities(
     X_str = X.astype('object')
     # With pandas
     vectorizer_cast.fit_transform(X_str)
-    check_same_transformers(expected_transformers_plain, vectorizer_cast.transformers)
+    check_same_transformers(expected_transformers_plain,
+                            vectorizer_cast.transformers)
     # With numpy
     vectorizer_cast.fit_transform(X_str.to_numpy())
-    check_same_transformers(expected_transformers_np_cast, vectorizer_cast.transformers)
+    check_same_transformers(expected_transformers_np_cast,
+                            vectorizer_cast.transformers)
 
 
 def test_with_clean_data():
@@ -210,7 +225,7 @@ def test_with_clean_data():
     )
 
 
-def test_with_dirty_data():
+def test_with_dirty_data() -> None:
     """
     Defines the expected returns of the vectorizer in different settings,
     and runs the tests with a dataset containing missing values.
@@ -255,7 +270,7 @@ def test_with_dirty_data():
     )
 
 
-def test_auto_cast():
+def test_auto_cast() -> None:
     """
     Tests that the SuperVectorizer automatic type detection works as expected.
     """
@@ -289,7 +304,8 @@ def test_auto_cast():
     X = _get_clean_dataframe()
     X_trans = vectorizer._auto_cast(X)
     for col in X_trans.columns:
-        assert type_equality(expected_types_clean_dataframe[col], X_trans[col].dtype)
+        assert type_equality(expected_types_clean_dataframe[col],
+                             X_trans[col].dtype)
 
     # Test that missing values don't prevent type detection
     expected_types_dirty_dataframe = {
@@ -304,11 +320,14 @@ def test_auto_cast():
     X = _get_dirty_dataframe()
     X_trans = vectorizer._auto_cast(X)
     for col in X_trans.columns:
-        assert type_equality(expected_types_dirty_dataframe[col], X_trans[col].dtype)
+        assert type_equality(expected_types_dirty_dataframe[col],
+                             X_trans[col].dtype)
+
 
 def test_with_arrays():
     """
-    Check that the SuperVectorizer works if we input a list of lists or a numpy array.
+    Check that the SuperVectorizer works if we input
+    a list of lists or a numpy array.
     """
     expected_transformers = {
         'numeric': [0, 1],
@@ -331,66 +350,67 @@ def test_with_arrays():
     check_same_transformers(expected_transformers, vectorizer.transformers)
 
 
-def test_get_feature_names_out():
+def test_get_feature_names_out() -> None:
     X = _get_clean_dataframe()
 
-    vectorizer_w_pass = SuperVectorizer(remainder='passthrough')
-    vectorizer_w_pass.fit(X)
+    vec_w_pass = SuperVectorizer(remainder='passthrough')
+    vec_w_pass.fit(X)
 
     if Version(sklearn.__version__) < Version('0.23'):
         with pytest.raises(NotImplementedError):
             # Prior to sklearn 0.23, ColumnTransformer.get_feature_names
             # with "passthrough" transformer(s) raises a NotImplementedError
-            assert vectorizer_w_pass.get_feature_names_out()
+            assert vec_w_pass.get_feature_names_out()
     else:
-        expected_feature_names_pass = [  # Order matters. If it doesn't, convert to set.
+        # In this test, order matters. If it doesn't, convert to set.
+        expected_feature_names_pass = [
             'str1_private', 'str1_public',
             'str2_chef', 'str2_lawyer', 'str2_manager', 'str2_officer', 'str2_teacher',
             'cat1_no', 'cat1_yes', 'cat2_20K+', 'cat2_30K+', 'cat2_40K+', 'cat2_50K+', 'cat2_60K+',
             'int', 'float'
         ]
         if Version(sklearn.__version__) < Version('1.0'):
-            assert vectorizer_w_pass.get_feature_names() == expected_feature_names_pass
+            assert vec_w_pass.get_feature_names() == expected_feature_names_pass
         else:
-            assert vectorizer_w_pass.get_feature_names_out() == expected_feature_names_pass
+            assert vec_w_pass.get_feature_names_out() == expected_feature_names_pass
 
-    vectorizer_w_drop = SuperVectorizer(remainder='drop')
-    vectorizer_w_drop.fit(X)
+    vec_w_drop = SuperVectorizer(remainder='drop')
+    vec_w_drop.fit(X)
 
-    expected_feature_names_drop = [  # Order matters. If it doesn't, convert to set.
+    # In this test, order matters. If it doesn't, convert to set.
+    expected_feature_names_drop = [
         'str1_private', 'str1_public',
         'str2_chef', 'str2_lawyer', 'str2_manager', 'str2_officer', 'str2_teacher',
         'cat1_no', 'cat1_yes', 'cat2_20K+', 'cat2_30K+', 'cat2_40K+', 'cat2_50K+', 'cat2_60K+'
     ]
     if Version(sklearn.__version__) < Version('1.0'):
-        assert vectorizer_w_drop.get_feature_names() == expected_feature_names_drop
+        assert vec_w_drop.get_feature_names() == expected_feature_names_drop
     else:
-        assert vectorizer_w_drop.get_feature_names_out() == expected_feature_names_drop
+        assert vec_w_drop.get_feature_names_out() == expected_feature_names_drop
 
 
-def test_fit():
+def test_fit() -> None:
     # Simply checks sklearn's `check_is_fitted` function raises an error if
     # the SuperVectorizer is instantiated but not fitted.
     # See GH#193
     sup_vec = SuperVectorizer()
     with pytest.raises(NotFittedError):
-        if Version(sklearn.__version__) >= Version('0.22'):
-            assert check_is_fitted(sup_vec)
-        else:
-            assert check_is_fitted(sup_vec, attributes=dir(sup_vec))
+        assert check_is_fitted(sup_vec)
 
 
-def test_transform():
+def test_transform() -> None:
     X = _get_clean_dataframe()
     sup_vec = SuperVectorizer()
     sup_vec.fit(X)
     s = [34, 5.5, 'private', 'manager', 'yes', '60K+']
     x = np.array(s).reshape(1, -1)
     x_trans = sup_vec.transform(x)
-    assert (x_trans == [[1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 34, 5.5]]).all()
+    assert (
+            x_trans == [[1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 34, 5.5]]
+    ).all()
 
 
-def test_fit_transform_equiv():
+def test_fit_transform_equiv() -> None:
     """
     We will test the equivalence between using `.fit_transform(X)`
     and `.fit(X).transform(X).`
@@ -417,24 +437,24 @@ def test_fit_transform_equiv():
 if __name__ == '__main__':
     print('start test_super_vectorizer with clean df')
     test_with_clean_data()
-    print('test_super_vectorizer with clean df passed')
+    print('passed')
     print('start test_super_vectorizer with dirty df')
     test_with_dirty_data()
-    print('test_super_vectorizer with dirty df passed')
+    print('passed')
     print('start test_auto_cast')
     test_auto_cast()
-    print('test_auto_cast passed')
+    print('passed')
     print("start test_with_arrays")
     test_with_arrays()
-    print("test_with_arrays passed")
+    print("passed")
     print('start  test_get_feature_names_out')
     test_get_feature_names_out()
-    print(' test_get_feature_names_out passed')
+    print('passed')
     print('start test_fit')
     test_fit()
-    print('test_fit passed')
+    print('passed')
     print('start fit_transform_equiv')
     test_fit_transform_equiv()
-    print('fit_transform_equiv passed')
+    print('passed')
 
     print('Done')
diff --git a/dirty_cat/test/test_target_encoder.py b/dirty_cat/test/test_target_encoder.py
index e3fbeb0..222b185 100644
--- a/dirty_cat/test/test_target_encoder.py
+++ b/dirty_cat/test/test_target_encoder.py
@@ -92,25 +92,32 @@ def test_target_encoder():
     encoder.fit(X, y)
 
     Ey_ = {0: 4/8, 1: 3/8, 2: 1/8}
-    Eyx_ = {0: {}, 1: {}, 2: {}}
-    Eyx_[0] = {'color': {'Red': 0,
-                         'red': 1,
-                         'green': 1/3,
-                         'blue': .5},
-               'gender': {'male': .5,
-                          'female': .5}}
-    Eyx_[1] = {'color': {'Red': 1,
-                         'red': 0,
-                         'green': 1/3,
-                         'blue': .5},
-               'gender': {'male': .5,
-                          'female': .25}}
-    Eyx_[2] = {'color': {'Red': 0,
-                         'red': 0,
-                         'green': 1/3,
-                         'blue': 0},
-               'gender': {'male': 0,
-                          'female': .25}}
+    Eyx_ = {
+        0: {
+            'color': {'Red': 0,
+                      'red': 1,
+                      'green': 1 / 3,
+                      'blue': .5},
+            'gender': {'male': .5,
+                       'female': .5},
+        },
+        1: {
+            'color': {'Red': 1,
+                      'red': 0,
+                      'green': 1 / 3,
+                      'blue': .5},
+            'gender': {'male': .5,
+                       'female': .25},
+        },
+        2: {
+            'color': {'Red': 0,
+                      'red': 0,
+                      'green': 1 / 3,
+                      'blue': 0},
+            'gender': {'male': 0,
+                       'female': .25},
+        },
+    }
     assert np.array_equal(np.unique(y), encoder.classes_)
     for k in [0, 1, 2]:
         assert Ey_[k] == encoder.Ey_[k]
@@ -128,15 +135,12 @@ def test_target_encoder():
 
     ans_dict = {0: {}, 1: {}, 2: {}}
     for k in [0, 1, 2]:
-        ans_dict[k] = {var:
-                       {cat:
-                        Eyx_[k][var][cat] *
-                        lambda_(count_[var][cat], n/len(count_[var])) +
-                        Ey_[k] *
-                        (1 - lambda_(count_[var][cat], n/len(count_[var])))
-                        for cat in Eyx_[k][var]}
-                       for var in Eyx_[k]
-                       }
+        ans_dict[k] = {
+            var: {
+                cat: Eyx_[k][var][cat] * lambda_(count_[var][cat], n/len(count_[var])) + Ey_[k] * (1 - lambda_(count_[var][cat], n/len(count_[var])))
+                for cat in Eyx_[k][var]}
+            for var in Eyx_[k]
+        }
 
     ans = np.zeros((n, 2*3))
     Xtest1 = np.array(['Red',
@@ -232,9 +236,7 @@ def _test_missing_values(input_type, missing):
 
     encoder = target_encoder.TargetEncoder(handle_missing=missing)
     if missing == 'error':
-        with pytest.raises(ValueError, match=r"Found missing values in input "
-                           "data; set handle_missing='' to encode "
-                           "with missing values"):
+        with pytest.raises(ValueError, match=r"Found missing values in input"):
             encoder.fit_transform(X, y)
         return
     elif missing == '':
@@ -248,8 +250,7 @@ def _test_missing_values(input_type, missing):
         assert dict(encoder.counter_[0]) == count_['color']
         assert dict(encoder.counter_[1]) == count_['gender']
     else:
-        with pytest.raises(ValueError, match=r"handle_missing"
-                           " should be either 'error' or ''"):
+        with pytest.raises(ValueError, match=r"expected any of"):
             encoder.fit_transform(X, y)
         return
 
@@ -305,9 +306,7 @@ def _test_missing_values_transform(input_type, missing):
                                            handle_missing=missing)
     if missing == 'error':
         encoder.fit_transform(X, y)
-        with pytest.raises(ValueError, match=r"Found missing values in input "
-                           "data; set handle_missing='' to encode "
-                           "with missing values"):
+        with pytest.raises(ValueError, match=r"Found missing values in input"):
             encoder.transform(X_test)
         return
     elif missing == '':
diff --git a/dirty_cat/test/test_utils.py b/dirty_cat/test/test_utils.py
index a339c34..e5e4293 100644
--- a/dirty_cat/test/test_utils.py
+++ b/dirty_cat/test/test_utils.py
@@ -7,11 +7,11 @@ def test_lrudict():
     dict_ = LRUDict(10)
 
     for x in range(15):
-        dict_[x] = 'filled'+str(x)
+        dict_[x] = f'filled {x}'
 
     for x in range(5, 15):
         assert x in dict_
-        assert dict_[x] == 'filled'+str(x)
+        assert dict_[x] == f'filled {x}'
 
     for x in range(5):
         assert (x not in dict_)