From 0d574505c13a9fa017f3e8d78696c0897cee2e89 Mon Sep 17 00:00:00 2001
From: Matt Davis <matteius@gmail.com>
Date: Sun, 11 Sep 2022 05:28:43 -0400
Subject: [PATCH] Allow adding packages to Pipfile in other sections.

---
 pipenv/cli/command.py |  3 ++-
 pipenv/core.py        | 17 +++++++++++++++--
 pipenv/project.py     | 11 +++++------
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/pipenv/cli/command.py b/pipenv/cli/command.py
index 306f34b3..d5882d8a 100644
--- a/pipenv/cli/command.py
+++ b/pipenv/cli/command.py
@@ -229,7 +229,7 @@ def cli(
 @skip_lock_option
 @install_options
 @pass_state
-def install(state, categories=None, **kwargs):
+def install(state, **kwargs):
     """Installs provided packages and adds them to Pipfile, or (if no packages are given), installs all packages from Pipfile."""
     from ..core import do_install
 
@@ -254,6 +254,7 @@ def install(state, categories=None, **kwargs):
         editable_packages=state.installstate.editables,
         site_packages=state.site_packages,
         extra_pip_args=state.installstate.extra_pip_args,
+        categories=state.installstate.categories,
     )
 
 
diff --git a/pipenv/core.py b/pipenv/core.py
index 10675717..ecf8c184 100644
--- a/pipenv/core.py
+++ b/pipenv/core.py
@@ -2039,6 +2039,7 @@ def do_install(
     selective_upgrade=False,
     site_packages=None,
     extra_pip_args=None,
+    categories=None,
 ):
     requirements_directory = vistir.path.create_tracked_tempdir(
         suffix="-requirements", prefix="pipenv-"
@@ -2313,13 +2314,21 @@ def do_install(
                             click.style("$ pipenv lock", fg="yellow"),
                         )
                     )
+                if categories:
+                    pipfile_sections = ""
+                    for c in categories:
+                        pipfile_sections += f"[{c}]"
+                elif dev:
+                    pipfile_sections = "[dev-packages]"
+                else:
+                    pipfile_sections = "[packages]"
                 sp.write(
                     "{} {} {} {}{}".format(
                         click.style("Adding", bold=True),
                         click.style(f"{pkg_requirement.name}", fg="green", bold=True),
                         click.style("to Pipfile's", bold=True),
                         click.style(
-                            "[dev-packages]" if dev else "[packages]",
+                            pipfile_sections,
                             fg="yellow",
                             bold=True,
                         ),
@@ -2333,7 +2342,11 @@ def do_install(
                     )
                     pkg_requirement.index = index_name
                 try:
-                    project.add_package_to_pipfile(pkg_requirement, dev)
+                    if categories:
+                        for category in categories:
+                            project.add_package_to_pipfile(pkg_requirement, dev, category)
+                    else:
+                        project.add_package_to_pipfile(pkg_requirement, dev)
                 except ValueError:
                     import traceback
 
diff --git a/pipenv/project.py b/pipenv/project.py
index 57183ee6..2f72b889 100644
--- a/pipenv/project.py
+++ b/pipenv/project.py
@@ -870,10 +870,9 @@ class Project:
 
     def get_package_name_in_pipfile(self, package_name, category=None, dev=False):
         """Get the equivalent package name in pipfile"""
-        if category:
-            section = self.parsed_pipfile.get(category, {})
-        else:  # TODO this branch will go away
+        if not category:
             category = "dev-packages" if dev else "packages"
+        section = self.parsed_pipfile.get(category, {})
         package_name = pep423_name(package_name)
         for name in section.keys():
             if pep423_name(name) == package_name:
@@ -905,7 +904,7 @@ class Project:
                 del parsed[section][pkg_name]
         self.write_toml(parsed)
 
-    def add_package_to_pipfile(self, package, dev=False):
+    def add_package_to_pipfile(self, package, dev=False, category=None):
         from .vendor.requirementslib import Requirement
 
         # Read and append Pipfile.
@@ -914,11 +913,11 @@ class Project:
         if not isinstance(package, Requirement):
             package = Requirement.from_line(package.strip())
         req_name, converted = package.pipfile_entry
-        key = "dev-packages" if dev else "packages"
+        key = category if category else "dev-packages" if dev else "packages"
         # Set empty group if it doesn't exist yet.
         if key not in p:
             p[key] = {}
-        name = self.get_package_name_in_pipfile(req_name, dev)
+        name = self.get_package_name_in_pipfile(req_name, category=category, dev=dev)
         if name and is_star(converted):
             # Skip for wildcard version
             return