1# Copyright 2023 The Bazel Authors. All rights reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"" 16 17load("//python/private:text_util.bzl", "render") 18load(":render_pkg_aliases.bzl", "render_multiplatform_pkg_aliases") 19load(":whl_config_setting.bzl", "whl_config_setting") 20 21_BUILD_FILE_CONTENTS = """\ 22package(default_visibility = ["//visibility:public"]) 23 24# Ensure the `requirements.bzl` source can be accessed by stardoc, since users load() from it 25exports_files(["requirements.bzl"]) 26""" 27 28def _impl(rctx): 29 bzl_packages = rctx.attr.packages or rctx.attr.whl_map.keys() 30 aliases = render_multiplatform_pkg_aliases( 31 aliases = { 32 key: _whl_config_settings_from_json(values) 33 for key, values in rctx.attr.whl_map.items() 34 }, 35 extra_hub_aliases = rctx.attr.extra_hub_aliases, 36 requirement_cycles = rctx.attr.groups, 37 ) 38 for path, contents in aliases.items(): 39 rctx.file(path, contents) 40 41 # NOTE: we are using the canonical name with the double '@' in order to 42 # always uniquely identify a repository, as the labels are being passed as 43 # a string and the resolution of the label happens at the call-site of the 44 # `requirement`, et al. macros. 45 macro_tmpl = "@@{name}//{{}}:{{}}".format(name = rctx.attr.name) 46 47 rctx.file("BUILD.bazel", _BUILD_FILE_CONTENTS) 48 rctx.template("requirements.bzl", rctx.attr._template, substitutions = { 49 "%%ALL_DATA_REQUIREMENTS%%": render.list([ 50 macro_tmpl.format(p, "data") 51 for p in bzl_packages 52 ]), 53 "%%ALL_REQUIREMENTS%%": render.list([ 54 macro_tmpl.format(p, "pkg") 55 for p in bzl_packages 56 ]), 57 "%%ALL_WHL_REQUIREMENTS_BY_PACKAGE%%": render.dict({ 58 p: macro_tmpl.format(p, "whl") 59 for p in bzl_packages 60 }), 61 "%%MACRO_TMPL%%": macro_tmpl, 62 }) 63 64hub_repository = repository_rule( 65 attrs = { 66 "extra_hub_aliases": attr.string_list_dict( 67 doc = "Extra aliases to make for specific wheels in the hub repo.", 68 mandatory = True, 69 ), 70 "groups": attr.string_list_dict( 71 mandatory = False, 72 ), 73 "packages": attr.string_list( 74 mandatory = False, 75 doc = """\ 76The list of packages that will be exposed via all_*requirements macros. Defaults to whl_map keys. 77""", 78 ), 79 "repo_name": attr.string( 80 mandatory = True, 81 doc = "The apparent name of the repo. This is needed because in bzlmod, the name attribute becomes the canonical name.", 82 ), 83 "whl_map": attr.string_dict( 84 mandatory = True, 85 doc = """\ 86The wheel map where values are json.encoded strings of the whl_map constructed 87in the pip.parse tag class. 88""", 89 ), 90 "_template": attr.label( 91 default = ":requirements.bzl.tmpl.bzlmod", 92 ), 93 }, 94 doc = """A rule for bzlmod mulitple pip repository creation. PRIVATE USE ONLY.""", 95 implementation = _impl, 96) 97 98def _whl_config_settings_from_json(repo_mapping_json): 99 """Deserialize the serialized values with whl_config_settings_to_json. 100 101 Args: 102 repo_mapping_json: {type}`str` 103 104 Returns: 105 What `whl_config_settings_to_json` accepts. 106 """ 107 return { 108 whl_config_setting(**v): repo 109 for repo, values in json.decode(repo_mapping_json).items() 110 for v in values 111 } 112 113def whl_config_settings_to_json(repo_mapping): 114 """A function to serialize the aliases so that `hub_repository` can accept them. 115 116 Args: 117 repo_mapping: {type}`dict[str, list[struct]]` repo to 118 {obj}`whl_config_setting` mapping. 119 120 Returns: 121 A deserializable JSON string 122 """ 123 return json.encode({ 124 repo: [_whl_config_setting_dict(s) for s in settings] 125 for repo, settings in repo_mapping.items() 126 }) 127 128def _whl_config_setting_dict(a): 129 ret = {} 130 if a.config_setting: 131 ret["config_setting"] = a.config_setting 132 if a.filename: 133 ret["filename"] = a.filename 134 if a.target_platforms: 135 ret["target_platforms"] = a.target_platforms 136 if a.version: 137 ret["version"] = a.version 138 return ret 139