From 354e1d15ada0984e4a26c6befff06e5974c7f909 Mon Sep 17 00:00:00 2001
From: Nick Ripley <nick.ripley@datadoghq.com>
Date: Wed, 11 Jun 2025 12:22:16 -0400
Subject: [PATCH] fix(profiling): make explicitly marked main packages "my
 code"

If a user runs their code as a "main module" which they installed like a
library (`python -m my_program args...`), and expicilty specifies that
it's the main package via DD_MAIN_PACKAGE, we should consider code from
that module "my code". Right now we don't, though, because we mark that
code as a "library" in the profiling code provenance info.  Check
whether the user specified a main package and make sure it gets marked
as "my code" by leaving the kind blank.
---
 .../internal/datadog/profiling/code_provenance.py  | 14 +++++++++++---
 tests/profiling_v2/test_code_provenance.py         |  5 +++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/ddtrace/internal/datadog/profiling/code_provenance.py b/ddtrace/internal/datadog/profiling/code_provenance.py
index b1c3a24ba75..2dc4cdd25e4 100644
--- a/ddtrace/internal/datadog/profiling/code_provenance.py
+++ b/ddtrace/internal/datadog/profiling/code_provenance.py
@@ -6,11 +6,11 @@
 import sysconfig
 import typing as t
 
+from ddtrace.internal import gitmetadata
 from ddtrace.internal.packages import _package_for_root_module_mapping
 
 
 class Library:
-
     def __init__(
         self,
         kind: str,
@@ -60,7 +60,7 @@ def __init__(self):
                     spec = importlib.util.find_spec(name)
                     if spec and spec.origin == "frozen":
                         python_stdlib.paths.add(f"<frozen {spec.name}>")
-                except Exception: # nosec
+                except Exception:  # nosec
                     continue
 
         self.libraries.append(python_stdlib)
@@ -69,6 +69,8 @@ def __init__(self):
 
         libraries: t.Dict[str, Library] = {}
 
+        _, _, main_package = gitmetadata.get_git_tags()
+
         site_packages = Path(sysconfig.get_path("purelib"))
         for module, dist in module_to_distribution.items():
             name = dist.name
@@ -76,9 +78,15 @@ def __init__(self):
             if module.startswith("__pycache__/"):
                 module = module[len("__pycache__/") :].split(".")[0] + ".py"
 
+            # If the user installed their code like a library and is running it
+            # as the main package (python -m my_package), and they explicitly
+            # specified that that's the main package, make sure it shows up as
+            # "my code" in the UI. Do this by leaving the "kind" blank
+            kind = "library" if name != main_package else ""
+
             lib = libraries.get(name)
             if lib is None:
-                lib = Library(kind="library", name=name, version=dist.version, paths=set())
+                lib = Library(kind=kind, name=name, version=dist.version, paths=set())
                 libraries[name] = lib
 
             # We assume that each module is a directory or a python file
diff --git a/tests/profiling_v2/test_code_provenance.py b/tests/profiling_v2/test_code_provenance.py
index fb4a9f30baf..3a13e59ee63 100644
--- a/tests/profiling_v2/test_code_provenance.py
+++ b/tests/profiling_v2/test_code_provenance.py
@@ -116,3 +116,8 @@ def test_stdlib_paths(self):
 
         for path in stdlib_paths:
             assert path.startswith("<frozen") or path == sysconfig.get_path("stdlib")
+
+    def test_main_package_my_code(self):
+        # TODO: set DD_MAIN_PACKAGE to something and confirm that the package
+        # does not have a "kind" in the code provenance
+        pass