Add simple wrapper to register nodes

sphinx/application.py

@@ -21,7 +21,6 @@ from inspect import isclass
 from os import path
 from typing import TYPE_CHECKING
 
-from docutils import nodes
 from docutils.parsers.rst import Directive, directives, roles
 from six import itervalues
 from six.moves import cStringIO
@@ -37,6 +36,7 @@ from sphinx.errors import (
 from sphinx.events import EventManager
 from sphinx.locale import __
 from sphinx.registry import SphinxComponentRegistry
+from sphinx.util import docutils
 from sphinx.util import import_object
 from sphinx.util import logging
 from sphinx.util import pycompat  # noqa: F401
@@ -48,6 +48,7 @@ from sphinx.util.tags import Tags
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Dict, IO, Iterable, Iterator, List, Tuple, Type, Union  # NOQA
+    from docutils import nodes  # NOQA
     from docutils.parsers import Parser  # NOQA
     from docutils.transform import Transform  # NOQA
     from sphinx.builders import Builder  # NOQA
@@ -611,13 +612,12 @@ class Sphinx(object):
            Added the support for keyword arguments giving visit functions.
         """
         logger.debug('[app] adding node: %r', (node, kwds))
-        if not kwds.pop('override', False) and \
-           hasattr(nodes.GenericNodeVisitor, 'visit_' + node.__name__):
+        if not kwds.pop('override', False) and docutils.is_node_registered(node):
             logger.warning(__('while setting up extension %s: node class %r is '
                               'already registered, its visitors will be overridden'),
                            self._setting_up_extension, node.__name__,
                            type='app', subtype='add_node')
-        nodes._add_node_class_names([node.__name__])
+        docutils.register_node(node)
         self.registry.add_translation_handlers(node, **kwds)
 
     def add_enumerable_node(self, node, figtype, title_getter=None, **kwds):

sphinx/util/docutils.py

@@ -34,13 +34,14 @@ logger = logging.getLogger(__name__)
 report_re = re.compile('^(.+?:(?:\\d+)?): \\((DEBUG|INFO|WARNING|ERROR|SEVERE)/(\\d+)?\\) ')
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Generator, Iterator, List, Tuple  # NOQA
+    from typing import Any, Callable, Generator, Iterator, List, Set, Tuple  # NOQA
     from docutils.statemachine import State, ViewList  # NOQA
     from sphinx.environment import BuildEnvironment  # NOQA
     from sphinx.io import SphinxFileInput  # NOQA
 
 
 __version_info__ = tuple(LooseVersion(docutils.__version__).version)
+additional_nodes = set()  # type: Set[nodes.Node]
 
 
 @contextmanager
@@ -56,6 +57,41 @@ def docutils_namespace():
         directives._directives = _directives
         roles._roles = _roles
 
+        for node in list(additional_nodes):
+            unregister_node(node)
+            additional_nodes.discard(node)
+
+
+def is_node_registered(node):
+    # type: (nodes.Node) -> bool
+    """Check the *node* is already registered."""
+    return hasattr(nodes.GenericNodeVisitor, 'visit_' + node.__name__)
+
+
+def register_node(node):
+    # type: (nodes.Node) -> None
+    """Register a node to docutils.
+
+    This modifies global state of some visitors.  So it is better to use this
+    inside ``docutils_namespace()`` to prevent side-effects.
+    """
+    if not hasattr(nodes.GenericNodeVisitor, 'visit_' + node.__name__):
+        nodes._add_node_class_names([node.__name__])
+        additional_nodes.add(node)
+
+
+def unregister_node(node):
+    # type: (nodes.Node) -> None
+    """Unregister a node from docutils.
+
+    This is inverse of ``nodes._add_nodes_class_names()``.
+    """
+    if hasattr(nodes.GenericNodeVisitor, 'visit_' + node.__name__):
+        delattr(nodes.GenericNodeVisitor, "visit_" + node.__name__)
+        delattr(nodes.GenericNodeVisitor, "depart_" + node.__name__)
+        delattr(nodes.SparseNodeVisitor, 'visit_' + node.__name__)
+        delattr(nodes.SparseNodeVisitor, 'depart_' + node.__name__)
+
 
 def patched_get_language(language_code, reporter=None):
     # type: (unicode, Reporter) -> Any

tests/test_util_docutils.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+"""
+    test_util_docutils
+    ~~~~~~~~~~~~~~~~~~
+
+    Tests util.utils functions.
+
+    :copyright: Copyright 2007-2018 by the Sphinx team, see AUTHORS.
+    :license: BSD, see LICENSE for details.
+"""
+
+from docutils import nodes
+
+from sphinx.util.docutils import docutils_namespace, register_node
+
+
+def test_register_node():
+    class custom_node(nodes.Element):
+        pass
+
+    with docutils_namespace():
+        register_node(custom_node)
+
+        # check registered
+        assert hasattr(nodes.GenericNodeVisitor, 'visit_custom_node')
+        assert hasattr(nodes.GenericNodeVisitor, 'depart_custom_node')
+        assert hasattr(nodes.SparseNodeVisitor, 'visit_custom_node')
+        assert hasattr(nodes.SparseNodeVisitor, 'depart_custom_node')
+
+    # check unregistered outside namespace
+    assert not hasattr(nodes.GenericNodeVisitor, 'visit_custom_node')
+    assert not hasattr(nodes.GenericNodeVisitor, 'depart_custom_node')
+    assert not hasattr(nodes.SparseNodeVisitor, 'visit_custom_node')
+    assert not hasattr(nodes.SparseNodeVisitor, 'depart_custom_node')