Quick attempt at making Sphinx mark classes with "final" when decorated with ``@_pytest.compat.final``.

pytest_final.py

"""
Quick attempt at making Sphinx mark classes with "final" when decorated with  ``@_pytest.compat.final``.

Ref: https://github.com/pytest-dev/pytest/pull/7780

-----

Based on Sphinx

Copyright (c) 2007-2020 by the Sphinx team.
All rights reserved.

See https://github.com/sphinx-doc/sphinx/blob/3.x/LICENSE
and https://github.com/sphinx-doc/sphinx/blob/3.x/AUTHORS
"""

# stdlib
import ast
from typing import Any, Dict, List

# 3rd party
import sphinx.pycode.ast
import sphinx.pycode.parser
from sphinx.application import Sphinx


def visit_Import(self, node: ast.Import) -> None:
	"""Handles Import node and record it to definition orders."""
	for name in node.names:
		self.add_entry(name.asname or name.name)

		if name.name == 'typing':
			self.typing = name.asname or name.name
		elif name.name == 'typing.final':
			self.typing_final = name.asname or name.name
		elif name.name == 'typing.overload':
			self.typing_overload = name.asname or name.name
		elif name.name == '_pytest.compat.final':
			self.pytest_final = name.asname or name.name


def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
	"""Handles Import node and record it to definition orders."""
	for name in node.names:
		self.add_entry(name.asname or name.name)

		if node.module == 'typing' and name.name == 'final':
			self.typing_final = name.asname or name.name
		elif node.module == 'typing' and name.name == 'overload':
			self.typing_overload = name.asname or name.name
		elif node.module == '_pytest.compat' and name.name == 'final':
			self.pytest_final = name.asname or name.name


def is_final(self, decorators: List[ast.expr]) -> bool:
	final = []
	if self.typing:
		final.append('%s.final' % self.typing)
	if self.typing_final:
		final.append(self.typing_final)
	if self.pytest_final:
		final.append(self.pytest_final)

	for decorator in decorators:
		try:
			if sphinx.pycode.ast.unparse(decorator) in final:
				return True
		except NotImplementedError:
			pass

	return False


def setup(app: Sphinx) -> Dict[str, Any]:
	sphinx.pycode.parser.VariableCommentPicker.is_final = is_final
	sphinx.pycode.parser.VariableCommentPicker.visit_Import = visit_Import
	sphinx.pycode.parser.VariableCommentPicker.visit_ImportFrom = visit_ImportFrom
	sphinx.pycode.parser.VariableCommentPicker.pytest_final = None

	return {}