From a42e4a5254bf2de2b2101a881fa5f3d40f4b2b3f Mon Sep 17 00:00:00 2001
From: Connor Baker <ConnorBaker01@gmail.com>
Date: Tue, 21 Oct 2025 04:12:41 +0000
Subject: [PATCH] python3Packages.tensorrt: fixup

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>
---
 .../python-modules/tensorrt/default.nix       | 40 ++++++++-----------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/pkgs/development/python-modules/tensorrt/default.nix b/pkgs/development/python-modules/tensorrt/default.nix
index 94a3504b95fa..86580550c7af 100644
--- a/pkgs/development/python-modules/tensorrt/default.nix
+++ b/pkgs/development/python-modules/tensorrt/default.nix
@@ -1,41 +1,35 @@
 {
+  autoPatchelfHook,
+  buildPythonPackage,
+  cudaPackages,
   lib,
   python,
-  autoAddDriverRunpath,
-  buildPythonPackage,
-  autoPatchelfHook,
-  unzip,
-  cudaPackages,
+  stdenv,
 }:
-
 let
-  pyVersion = "${lib.versions.major python.version}${lib.versions.minor python.version}";
-  buildVersion = lib.optionalString (cudaPackages ? tensorrt) cudaPackages.tensorrt.version;
+  inherit (cudaPackages.tensorrt) src pname version;
+  inherit (lib.versions) major minor;
+  inherit (stdenv.hostPlatform) parsed;
 in
 buildPythonPackage {
-  pname = "tensorrt";
-  version = buildVersion;
+  inherit pname version;
 
-  src = cudaPackages.tensorrt.src;
+  src =
+    let
+      # https://peps.python.org/pep-0427/#file-name-convention
+      distribution = pname;
+      pythonTag = "cp${major python.version}${minor python.version}";
+      abiTag = "none";
+      platformTag = "${parsed.kernel.name}_${parsed.cpu.name}";
+    in
+    src + "/python/${distribution}-${version}-${pythonTag}-${abiTag}-${platformTag}.whl";
 
   format = "wheel";
-  # We unpack the wheel ourselves because of the odd packaging.
-  dontUseWheelUnpack = true;
 
   nativeBuildInputs = [
-    unzip
     autoPatchelfHook
-    autoAddDriverRunpath
   ];
 
-  preUnpack = ''
-    mkdir -p dist
-    tar --strip-components=2 -xf "$src" --directory=dist \
-      "TensorRT-${buildVersion}/python/tensorrt-${buildVersion}-cp${pyVersion}-none-linux_x86_64.whl"
-  '';
-
-  sourceRoot = ".";
-
   buildInputs = [
     cudaPackages.cudnn
     cudaPackages.tensorrt