pytorch · dagitses · Mar 21, 2022 · Mar 21, 2022 · Mar 21, 2022 · Mar 21, 2022
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -31,7 +31,7 @@ py_binary(
     name = "gen",
     srcs = ["tools/setup_helpers/gen.py"],
     deps = [
-        ":tools_codegen"
+        "//tools/codegen"
     ],
 )
 
@@ -105,19 +105,14 @@ generate_aten(
     generator=":gen",
 )
 
-py_library(
-    name = "tools_codegen",
-    srcs = glob(["tools/codegen/**/*.py"]),
-)
-
 py_library(
     name = "tools_autograd",
     srcs = glob(["tools/autograd/*.py"]),
     data = glob([
         "tools/autograd/*.yaml",
         "tools/autograd/templates/*",
     ]),
-    deps = [":tools_codegen"],
+    deps = ["//tools/codegen"],
 )
 
 py_library(
@@ -1380,7 +1375,7 @@ cc_library(
 py_binary(
     name = "gen_op",
     srcs = ["caffe2/contrib/aten/gen_op.py"],
-    deps = [":tools_codegen"],
+    deps = ["//tools/codegen"],
 )
 
 genrule(

diff --git a/tools/bazel.bzl b/tools/bazel.bzl
@@ -3,6 +3,13 @@ load("@rules_cuda//cuda:defs.bzl", "requires_cuda_enabled")
 load("//c10/macros:cmake_configure_file.bzl", "cmake_configure_file")
 load("//tools/config:defs.bzl", "if_cuda")
 
+def _py_library(name, **kwds):
+    deps = [dep for dep in kwds.pop("deps", []) if dep != None]
+    native.py_library(name = name, deps = deps, **kwds)
+
+def _requirement(_pypi_project):
+    return None
+
 # Rules implementation for the Bazel build system. Since the common
 # build structure aims to replicate Bazel as much as possible, most of
 # the rules simply forward to the Bazel definitions.
@@ -14,6 +21,8 @@ rules = struct(
     filegroup = native.filegroup,
     glob = native.glob,
     if_cuda = if_cuda,
+    py_library = _py_library,
+    requirement = _requirement,
     requires_cuda_enabled = requires_cuda_enabled,
     select = select,
     test_suite = native.test_suite,

diff --git a/tools/codegen/BUILD.bazel b/tools/codegen/BUILD.bazel
@@ -0,0 +1,4 @@
+load("//:tools/bazel.bzl", "rules")
+load(":build.bzl", "define_targets")
+
+define_targets(rules = rules)