diff --git a/docs/rules.md b/docs/rules.md
index cfb65f19..88cd1e65 100644
--- a/docs/rules.md
+++ b/docs/rules.md
@@ -105,7 +105,7 @@ py_pytest_main wraps the template rendering target and the final py_library.
 | Name  | Description | Default Value |
 | :------------- | :------------- | :------------- |
 | <a id="py_pytest_main-name"></a>name |  The name of the runable target that updates the test entry file.   |  none |
-| <a id="py_pytest_main-py_library"></a>py_library |  Use this attribute to override the default py_library rule.   |  <code>&lt;function py_library&gt;</code> |
+| <a id="py_pytest_main-py_library"></a>py_library |  Use this attribute to override the default py_library rule.   |  <code>&lt;unknown object com.google.devtools.build.skydoc.fakebuildapi.FakeStarlarkRuleFunctionsApi$RuleDefinitionIdentifier&gt;</code> |
 | <a id="py_pytest_main-deps"></a>deps |  A list containing the pytest library target, e.g., @pypi_pytest//:pkg.   |  <code>[]</code> |
 | <a id="py_pytest_main-data"></a>data |  A list of data dependencies to pass to the py_library target.   |  <code>[]</code> |
 | <a id="py_pytest_main-testonly"></a>testonly |  A boolean indicating if the py_library target is testonly.   |  <code>True</code> |
diff --git a/examples/pytest/BUILD.bazel b/examples/pytest/BUILD.bazel
index aea9922d..b02fe563 100644
--- a/examples/pytest/BUILD.bazel
+++ b/examples/pytest/BUILD.bazel
@@ -1,5 +1,4 @@
-load("@rules_python//python:defs.bzl", "py_test")
-load("@//py:defs.bzl", "py_pytest_main")
+load("@aspect_rules_py//py:defs.bzl", "py_pytest_main", "py_test")
 
 py_pytest_main(
     name = "__test__",
diff --git a/py/private/py_pytest_main.bzl b/py/private/py_pytest_main.bzl
index af9cbf58..3723bbf7 100644
--- a/py/private/py_pytest_main.bzl
+++ b/py/private/py_pytest_main.bzl
@@ -15,7 +15,7 @@
 """py_test entrypoint generation.
 """
 
-load("@rules_python//python:defs.bzl", default_py_library = "py_library")
+load(":py_library.bzl", default_py_library = "py_library")
 
 def _py_pytest_main_impl(ctx):
     substitutions = {