[Bugfix] add DGLWarning class (#1898)

* add DGLWarning class

* docstrings

* lint

python/dgl/base.py

@@ -26,12 +26,24 @@ def is_all(arg):
     """Return true if the argument is a special symbol for all nodes or edges."""
     return isinstance(arg, str) and arg == ALL
 
+# pylint: disable=invalid-name
+_default_formatwarning = warnings.formatwarning
+
+class DGLWarning(UserWarning):
+    """DGL Warning class."""
+
 # pylint: disable=unused-argument
-def dgl_warning_format(message, category, filename, lineno, file=None, line=None):
+def dgl_warning_format(message, category, filename, lineno, line=None):
     """Format DGL warnings."""
+    if isinstance(category, DGLWarning):
         return "DGL Warning: {}\n".format(message)
+    else:
+        return _default_formatwarning(message, category, filename, lineno, line=None)
+
+def dgl_warning(message, category=DGLWarning, stacklevel=1):
+    """DGL warning wrapper that defaults to ``DGLWarning`` instead of ``UserWarning`` category."""
+    return warnings.warn(message, category=category, stacklevel=1)
 
 warnings.formatwarning = dgl_warning_format
-dgl_warning = warnings.warn  # pylint: disable=invalid-name
 
 _init_internal_api()