From df3de222e8666c0a48b4dc4fd6dec69392307873 Mon Sep 17 00:00:00 2001
From: Itai Hay <3392524+itaihay@users.noreply.github.com>
Date: Mon, 27 Jan 2025 21:44:38 +0200
Subject: [PATCH] Add N808 for type variable name convention check

---
 README.rst            |  4 ++++
 src/pep8ext_naming.py | 23 +++++++++++++++++++++
 testsuite/N808.py     | 47 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+)
 create mode 100644 testsuite/N808.py

diff --git a/README.rst b/README.rst
index b016ef7..362b932 100644
--- a/README.rst
+++ b/README.rst
@@ -54,6 +54,9 @@ These error codes are emitted:
 +---------+-----------------------------------------------------------------+
 | _`N807` | function name should not start and end with '__'                |
 +---------+-----------------------------------------------------------------+
+| _`N808` | type variable names should use CapWords convention and an       |
+|         | optional suffix '_co' or '_contra' (`type variable names`_)     |
++---------+-----------------------------------------------------------------+
 | _`N811` | constant imported as non constant (`constants`_)                |
 +---------+-----------------------------------------------------------------+
 | _`N812` | lowercase imported as non-lowercase                             |
@@ -80,6 +83,7 @@ These error codes are emitted:
 .. _function names: https://www.python.org/dev/peps/pep-0008/#function-and-variable-names
 .. _function arguments: https://www.python.org/dev/peps/pep-0008/#function-and-method-arguments
 .. _method names: https://www.python.org/dev/peps/pep-0008/#method-names-and-instance-variables
+.. _type variable names: https://peps.python.org/pep-0008/#type-variable-names
 
 Options
 -------
diff --git a/src/pep8ext_naming.py b/src/pep8ext_naming.py
index aaa5548..0b713f0 100644
--- a/src/pep8ext_naming.py
+++ b/src/pep8ext_naming.py
@@ -254,6 +254,7 @@ class ClassNameCheck(BaseASTCheck):
     Classes for internal use have a leading underscore in addition.
     """
     N801 = "class name '{name}' should use CapWords convention"
+    N808 = "type variable name '{name}' should use CapWords convention and an optional suffix '_co' or '_contra'"
     N818 = "exception name '{name}' should be named with an Error suffix"
 
     @classmethod
@@ -275,6 +276,28 @@ def superclass_names(cls, name, parents: Iterable, _names=None):
                 names.update(cls.superclass_names(base.id, parents, names))
         return names
 
+    def visit_module(self, node, parents: Iterable, ignore=None):
+        for body in node.body:
+            try:
+                if hasattr(body, 'targets') and len(body.targets) == 1:
+                    name = body.targets[0].id
+                    func_name = body.value.func.id
+            except AttributeError:
+                continue
+
+            if _ignored(name, ignore) or func_name != "TypeVar":
+                return
+
+            if not name[:1].isupper():
+                yield self.err(body, 'N808', name=name)
+            elif '_' in name:
+                underscore_split = name.split('_')
+                suffix = underscore_split[-1]
+                if len(underscore_split) > 2 or (suffix != 'co' and suffix != 'contra'):
+                    yield self.err(body, 'N808', name=name)
+
+
+
     def visit_classdef(self, node, parents: Iterable, ignore=None):
         name = node.name
         if _ignored(name, ignore):
diff --git a/testsuite/N808.py b/testsuite/N808.py
new file mode 100644
index 0000000..61a49d1
--- /dev/null
+++ b/testsuite/N808.py
@@ -0,0 +1,47 @@
+from typing import TypeVar
+
+#: Okay
+Ok = TypeVar("Ok")
+
+#: N808
+notok = TypeVar("notok")
+
+#: N808
+notok_co = TypeVar("notok_co")
+
+#: Okay(--ignore-names=notok)
+notok = TypeVar("notok")
+
+#: N808:1:1(--ignore-names=*OK)
+notok = TypeVar("notok")
+
+#: Okay
+Ok_co = TypeVar("Ok_co", covariant=True)
+
+#: Okay
+Ok_contra = TypeVar("Ok_contra", contravariant=True)
+
+#: N808
+Ok__contra = TypeVar("Ok__contra", contravariant=True)
+
+#: Okay
+Ok_co = TypeVar("Ok_co")
+
+#: Okay
+Ok_contra = TypeVar("Ok_contra")
+
+#: Okay
+Good = TypeVar("Good")
+
+#: N808
+__NotGood = TypeVar("__NotGood")
+
+#: N808
+__NotGood__ = TypeVar("__NotGood__")
+
+#: N808
+NotGood__ = TypeVar("NotGood__")
+
+# Make sure other function calls do not get checked
+#: Okay
+t = str('something')