diff --git a/pysrc/setup_file.py b/pysrc/setup_file.py
index de3bc5de..2fe38585 100644
--- a/pysrc/setup_file.py
+++ b/pysrc/setup_file.py
@@ -10,7 +10,7 @@ def _save_passed_args(**kwargs):
 
     distutils.core.setup = _save_passed_args
     setuptools.setup = _save_passed_args
-    setup_py_content = setup_py_content.replace("setup(", "passed_arguments = setup(")
+    setup_py_content = re.sub(r"(setuptools\.)?setup\(", r"passed_arguments = \g<0>", setup_py_content)
 
     # Fetch the arguments that were passed to the setup.py
     exec(setup_py_content, globals())
diff --git a/test/unit/setup_file.spec.ts b/test/unit/setup_file.spec.ts
new file mode 100644
index 00000000..0eb88707
--- /dev/null
+++ b/test/unit/setup_file.spec.ts
@@ -0,0 +1,18 @@
+import path = require('path');
+import { executeSync } from '../../lib/dependencies/sub-process';
+
+describe('Test setup_file.py', () => {
+  it.each`
+    setupPyContent
+    ${'import setuptools;setuptools.setup(name="test")'}
+    ${'from setuptools import setup;setup(name="test")'}
+  `("parse works for '$setupPyContent'", ({ setupPyContent }) => {
+    const result = executeSync(
+      'python',
+      ['-c', `from setup_file import parse; parse('${setupPyContent}')`],
+      { cwd: path.resolve(__dirname, '../../pysrc') }
+    );
+
+    expect(result.status).toBe(0);
+  });
+});