def _find_tensorrt_config()

in gpu/find_cuda_config.py [0:0]

• 25 lines of code
• 4 McCabe index (conditional complexity)

def _find_tensorrt_config(base_paths, required_version):

  def get_header_version(path):
    version = (
        _get_header_version(path, name)
        for name in ("NV_TENSORRT_MAJOR", "NV_TENSORRT_MINOR",
                     "NV_TENSORRT_PATCH"))
    # `version` is a generator object, so we convert it to a list before using
    # it (muitiple times below).
    version = list(version)
    if not all(version):
      return None  # Versions not found, make _matches_version returns False.
    return ".".join(version)

  try:
    header_path, header_version = _find_header(base_paths, "NvInfer.h",
                                               required_version,
                                               get_header_version)
  except ConfigError:
    # TensorRT 6 moved the version information to NvInferVersion.h.
    header_path, header_version = _find_header(base_paths, "NvInferVersion.h",
                                               required_version,
                                               get_header_version)

  tensorrt_version = header_version.split(".")[0]
  library_path = _find_library(base_paths, "nvinfer", tensorrt_version)

  return {
      "tensorrt_version": tensorrt_version,
      "tensorrt_include_dir": os.path.dirname(header_path),
      "tensorrt_library_dir": os.path.dirname(library_path),
  }