def _get_policies_for_line_with_model_urls()

in tools/validator.py [0:0]

• 15 lines of code
• 3 McCabe index (conditional complexity)

def _get_policies_for_line_with_model_urls(
    line: str) -> Iterator[ModelParsingPolicy]:
  """Yields a parsing policy for models that correspond to found tfhub.dev URLs.

  For each tfhub.dev URL that can be found in the given string, the
  corresponding parsing policy is yielded. The URL must match
  https://tfhub.dev/PUBLISHER/((tfjs|lite|coral)-model/)NAME(/VERSION) so a
  PUBLISHER and a NAME must be specified while a VERSION is optional. If no
  version is specified, a '*' will be used to simplify finding all documentation
  files that could be rendered.

  Args:
    line: String that could contain a tfhub.dev model URL.

  Yields:
    A ModelParsingPolicy which can be used to find the documentation files for
    the model corresponding to the URL.
  """
  for url_match in _TFHUB_MODEL_URL_PATTERN.finditer(line):
    groupdict = url_match.groupdict()
    publisher = groupdict.get("publisher")
    prefix = groupdict.get("prefix")
    name = groupdict.get("name")

    trailing_version = re.search(rf"(?<=/){_MODEL_VERSION_PATTERN}$", name)
    if trailing_version is None:
      version = "*"
    else:
      # Move the trailing version from the model name to version.
      version = trailing_version.group(0)
      name = re.sub(f"/{version}$", "", name)

    policy_class = POLICY_BY_URL_PREFIX.get(prefix, SavedModelParsingPolicy)
    yield policy_class({}, publisher, name, version)