Implement class checking.

Author: domdfcoding

flake8_params/__init__.py

@@ -30,12 +30,12 @@
 
 # stdlib
 import ast
-from typing import Iterator, Union
+from typing import Iterator, List, Optional, Union
 
 # 3rd party
 import flake8_helper
 
-__all__ = ("Plugin", "Visitor", "get_decorator_names")
+__all__ = ("Plugin", "Visitor", "get_decorator_names", "check_params")
 
 __author__ = "Dominic Davis-Foster"
 __copyright__ = "2025 Dominic Davis-Foster"
@@ -46,6 +46,7 @@
 PRM001 = "PRM001 Docstring parameters in wrong order."
 PRM002 = "PRM002 Missing parameters in docstring."
 PRM003 = "PRM003 Extra parameters in docstring."
+# TODO: class-specific codes?
 
 deco_allowed_attr_names = {
 		".setter",  # Property setter
@@ -73,7 +74,7 @@ def _get_deco_name(decorator: ast.expr) -> Iterator[str]:
 		raise NotImplementedError(decorator)
 
 
-def get_decorator_names(function: Union[ast.AsyncFunctionDef, ast.FunctionDef]) -> Iterator[str]:
+def get_decorator_names(function: Union[ast.AsyncFunctionDef, ast.FunctionDef, ast.ClassDef]) -> Iterator[str]:
 	"""
 	Returns an iterator of the dotted names of decorators for the given function.
 
@@ -84,81 +85,136 @@ def get_decorator_names(function: Union[ast.AsyncFunctionDef, ast.FunctionDef])
 		yield from _get_deco_name(decorator)
 
 
+def check_params(
+		signature_args: List[str],
+		docstring_args: List[str],
+		decorators: List[str],
+		) -> Optional[str]:
+	"""
+	Check if signature and docstring parameters match, and return the flake8 error code if not.
+
+	:param signature_args:
+	:param docstring_args:
+	:param decorators: List of dotted names (e.g. ``foo.bar``, for ``@foo.bar()``) of decorators for the function or class.
+
+	:returns: Either a flake8 error code and description, or :py:obj:`None` if no errors were detected.
+	"""
+
+	if "self" in signature_args:
+		signature_args.remove("self")
+
+	if "classmethod" in decorators and signature_args:
+		signature_args.pop(0)
+	for deco in decorators:
+		if any(deco.endswith(name) for name in deco_allowed_attr_names):
+			signature_args = []
+			break
+
+	if not signature_args and not docstring_args:
+		# No args either way
+		return None
+
+	if signature_args == docstring_args:
+		# All match
+		return None
+
+	# Either wrong order, extra in signature, extra in doc
+	signature_set = set(signature_args)
+	docstring_set = set(docstring_args)
+	if signature_set == docstring_set:
+		# Wrong order
+		return PRM001
+	elif signature_set - docstring_set:
+		# Extras in signature
+		return PRM002
+	elif docstring_set - signature_set:
+		# Extras in docstrings
+		return PRM003
+
+	return None  # pragma: no cover
+
+
 class Visitor(flake8_helper.Visitor):
 	"""
 	AST node visitor for identifying mismatches between function signatures and docstring params.
 	"""
 
+	# TODO: async functions
+
 	def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # noqa: D102
-		docstring = ast.get_docstring(node, clean=False)
 		if node.name == "__init__":
-			# TODO: special case; parameters go on class
 			self.generic_visit(node)
 			return
 
+		docstring = ast.get_docstring(node, clean=False)
+
 		if not docstring:
 			self.generic_visit(node)
 			return
 
-		seen_args = []
+		docstring_args = []
 		for line in docstring.split('\n'):
 			line = line.strip()
 			if line.startswith(":param"):
-				seen_args.append(line[6:].split(':', 1)[0].strip())
+				docstring_args.append(line[6:].split(':', 1)[0].strip())
 
 		signature_args = [a.arg for a in node.args.args]
-		if "self" in signature_args:
-			signature_args.remove("self")
 
-		# decorators = [n.id for n in node.decorator_list if isinstance(n, ast.Name)]
 		decorators = list(get_decorator_names(node))
-		if "classmethod" in decorators and signature_args:
-			signature_args.pop(0)
-		for deco in decorators:
-			if any(deco.endswith(name) for name in deco_allowed_attr_names):
-				signature_args = []
-				break
 
-		if not signature_args and not seen_args:
-			# No args either way
+		error = check_params(signature_args, docstring_args, decorators)
+		if not error:
+			self.generic_visit(node)
+			return
+
+		self.errors.append((
+				node.lineno,
+				node.col_offset,
+				error,
+				))
+
+		self.generic_visit(node)
+
+	def visit_ClassDef(self, node: ast.ClassDef) -> None:  # noqa: D102
+		docstring = ast.get_docstring(node, clean=False)
+
+		if not docstring:
 			self.generic_visit(node)
 			return
 
-		if signature_args == seen_args:
-			# All match
+		docstring_args = []
+		for line in docstring.split('\n'):
+			line = line.strip()
+			if line.startswith(":param"):
+				docstring_args.append(line[6:].split(':', 1)[0].strip())
+
+		decorators = list(get_decorator_names(node))
+
+		signature_args = []
+		functions_in_body: List[ast.FunctionDef] = [n for n in node.body if isinstance(n, ast.FunctionDef)]
+
+		for function in functions_in_body:
+			if function.name == "__init__":
+				signature_args = [a.arg for a in function.args.args]
+				break
+		else:
+			# No __init__; maybe it comes from a base class.
+			# TODO: check for base classes and still error if non exist
+			return None
+
+		error = check_params(signature_args, docstring_args, decorators)
+		if not error:
 			self.generic_visit(node)
 			return
 
-		# Either wrong order, extra in signature, extra in doc
-		signature_set = set(signature_args)
-		seen_set = set(seen_args)
-		if signature_set == seen_set:
-			# Wrong order
-			self.errors.append((
-					node.lineno,
-					node.col_offset,
-					PRM001,
-					))
-		elif signature_set - seen_set:
-			# Extras in signature
-			self.errors.append((
-					node.lineno,
-					node.col_offset,
-					PRM002,
-					))
-		elif seen_set - signature_set:
-			# Extras in docstrings
-			self.errors.append((
-					node.lineno,
-					node.col_offset,
-					PRM003,
-					))
+		self.errors.append((
+				node.lineno,
+				node.col_offset,
+				error,
+				))
 
 		self.generic_visit(node)
 
-	# def visit_ClassDef(self, node: ast.ClassDef):
-	# 	breakpoint()
-
 
 class Plugin(flake8_helper.Plugin[Visitor]):
 	"""

tests/example_code.py

@@ -36,8 +36,7 @@ def docstring_wrong_order(self, foo, bar, baz):
 	"""
 
 
-@classmethod
-def is_classmethod(cls, foo, bar, baz):
+def missing_in_signature(foo, bar):
 	"""
 	Does something.
 
@@ -47,60 +46,135 @@ def is_classmethod(cls, foo, bar, baz):
 	"""
 
 
-@classmethod
-def missing_in_docstring_with_classmethod(cls, foo, bar, baz):
+def missing_in_signature_with_self(self, foo, bar):
 	"""
 	Does something.
 
 	:param foo:
+	:param bar:
 	:param baz:
 	"""
 
 
-def missing_in_signature(foo, bar):
+class ClassMissingDocstring:
 	"""
-	Does something.
+	A class.
 
 	:param foo:
 	:param bar:
-	:param baz:
 	"""
 
+	def __init__(self, foo, bar, baz):
+		pass
 
-def missing_in_signature_with_self(self, foo, bar):
+	@classmethod
+	def is_classmethod(cls, foo, bar, baz):
+		"""
+		Does something.
+
+		:param foo:
+		:param bar:
+		:param baz:
+		"""
+
+	@classmethod
+	def missing_in_docstring_with_classmethod(cls, foo, bar, baz):
+		"""
+		Does something.
+
+		:param foo:
+		:param baz:
+		"""
+
+	@classmethod
+	def missing_in_signature_with_classmethod(cls, foo, baz):
+		"""
+		Does something.
+
+		:param foo:
+		:param bar:
+		:param baz:
+		"""
+
+	@property
+	def a_property(self):
+		"""
+		A property.
+		"""
+
+	@a_property.setter
+	def a_property(self, val):
+		"""
+		A property.
+		"""
+
+
+class ClassMissingSignature:
 	"""
-	Does something.
+	A class.
 
 	:param foo:
 	:param bar:
 	:param baz:
 	"""
 
+	def __init__(self, foo, bar):
+		pass
 
-@classmethod
-def missing_in_signature_with_classmethod(cls, foo, baz):
+
+class ClassWrongOrder:
 	"""
-	Does something.
+	A class.
 
-	:param foo:
 	:param bar:
+	:param foo:
 	:param baz:
 	"""
 
+	def __init__(self, foo, bar, baz):
+		pass
+
 
-@property
-def a_property(self):
+class ClassNoInit:
 	"""
-	A property.
+	A class.
+
+	:param foo:
+	:param bar:
+	:param baz:
 	"""
 
 
-@a_property.setter
-def a_property(self, val):
+class ClassNoDocstring:
+
+	def __init__(self, foo, bar, baz):
+		pass
+
+
+class ClassInitDocstring:
+
+	def __init__(self, foo, bar, baz):
+		"""
+		Setup the function.
+
+		:param foo:
+		:param bar:
+		:param baz:
+		"""
+
+
+class GoodClass:
 	"""
-	A property.
+	A class.
+
+	:param foo:
+	:param bar:
+	:param baz:
 	"""
 
+	def __init__(self, foo, bar, baz):
+		pass
+
 
 # 3rd party
 import click  # type: ignore

tests/test_flake8_params_/test_plugin.txt

@@ -1,7 +1,10 @@
 1:0: PRM002 Missing parameters in docstring.
 20:0: PRM002 Missing parameters in docstring.
 29:0: PRM001 Docstring parameters in wrong order.
-51:0: PRM002 Missing parameters in docstring.
-60:0: PRM003 Extra parameters in docstring.
-70:0: PRM003 Extra parameters in docstring.
-81:0: PRM003 Extra parameters in docstring.
+39:0: PRM003 Extra parameters in docstring.
+49:0: PRM003 Extra parameters in docstring.
+59:0: PRM002 Missing parameters in docstring.
+81:1: PRM002 Missing parameters in docstring.
+90:1: PRM003 Extra parameters in docstring.
+112:0: PRM003 Extra parameters in docstring.
+125:0: PRM001 Docstring parameters in wrong order.

tox.ini

@@ -177,7 +177,7 @@ rst-directives =
 rst-roles = choosealicense
 per-file-ignores =
     tests/*: D100 D101 D102 D103 D104 D106 D201 D204 D207 D208 D209 D210 D211 D212 D213 D214 D215 D300 D301 D400 D402 D403 D404 D415 D417 DALL000  SLOT000 SLOT001 SLOT002
-    tests/example_code.py: PRM001 PRM002 PRM003 MAN001 MAN002 D100 D102 D103 DALL000 E402
+    tests/example_code.py: PRM001 PRM002 PRM003 MAN001 MAN002 D100 D101 D102 D103 DALL000 E402
     */*.pyi: E301 E302 E305 D100 D101 D102 D103 D104 D106 D201 D204 D207 D208 D209 D210 D211 D212 D213 D214 D215 D300 D301 D400 D402 D403 D404 D415 D417 DALL000  SLOT000 SLOT001 SLOT002
 pytest-parametrize-names-type = csv
 inline-quotes = "