Add source code links to documentation

Adjust page width CSS

docs/_static/css/custom.css

@@ -27,8 +27,10 @@
     background: #85cfcb;
 }
 
-.wy-nav-content {
-    max-width: 1000px;
+@media screen and (min-width: 1000px) {
+    .wy-nav-content {
+        max-width: 1000px;
+    }
 }
 
 html.writer-html4 .rst-content dl:not(.docutils) > dt, html.writer-html5 .rst-content dl[class]:not(.option-list):not(.field-list):not(.footnote):not(.citation):not(.glossary):not(.simple) > dt {
@@ -50,4 +52,8 @@ h1, h2, h3, h4, h5, h6 {
 
 div.toctree-wrapper .caption-text{
     color: #203847;
+}
+
+.rst-content .viewcode-back, .rst-content .viewcode-link {
+    padding-left: 6px;
 }

docs/conf.py

@@ -1,5 +1,8 @@
 import os
 import sys
+import importlib
+import inspect
+import pathlib
 
 sys.path.insert(0, os.path.abspath(".."))  # pylint: disable=wrong-import-position
 
@@ -22,6 +25,7 @@
     "sphinx_rtd_theme",
     "sphinx.ext.napoleon",
     "sphinx.ext.mathjax",
+    "sphinx.ext.linkcode",
     "myst_parser",
 ]
 
@@ -62,3 +66,24 @@
     "collapse_navigation": False,
     "prev_next_buttons_location": "None",
 }
+
+
+def linkcode_resolve(domain, info):
+    github_url = f"https://github.com/CQCL/qujax/tree/develop/qujax"
+
+    if domain != "py":
+        return
+
+    module = importlib.import_module(info["module"])
+    obj = getattr(module, info["fullname"])
+
+    try:
+        path = pathlib.Path(inspect.getsourcefile(obj))
+        file_name = path.name
+        lines = inspect.getsourcelines(obj)
+    except TypeError:
+        return
+
+    start, end = lines[1], lines[1] + len(lines[0]) - 1
+
+    return f"{github_url}/{file_name}#L{start}-L{end}"