[ADD] ORM: with_company

3 files changed, 100 insertions, 12 deletions

diff --git a/doc/reference/orm.rst b/doc/reference/orm.rst
index d6ada51b4d4..d9e9a715a57 100644
--- a/doc/reference/orm.rst
+++ b/doc/reference/orm.rst
@@ -593,6 +593,8 @@ Altering the environment
 
 .. automethod:: Model.with_user
 
+.. automethod:: Model.with_company
+
 .. automethod:: Model.with_env
 
 .. automethod:: Model.sudo
diff --git a/odoo/addons/test_new_api/tests/test_company_checks.py b/odoo/addons/test_new_api/tests/test_company_checks.py
index 94645fe1665..8d1aecaebfa 100644
--- a/odoo/addons/test_new_api/tests/test_company_checks.py
+++ b/odoo/addons/test_new_api/tests/test_company_checks.py
@@ -22,6 +22,16 @@ class TestCompanyCheck(common.TransactionCase):
             'name': 'M2',
             'company_id': self.company_b.id,
         })
+        self.company_c = self.env['res.company'].create({
+            'name': 'Company C'
+        })
+
+        self.test_user = self.env['res.users'].create({
+            'name': 'Test',
+            'login': 'test',
+            'company_id': self.company_a.id,
+            'company_ids': (self.company_a | self.company_c).ids,
+        })
 
     def test_company_check_0(self):
         """ Check the option _check_company_auto is well set on records"""
@@ -75,18 +85,7 @@ class TestCompanyCheck(common.TransactionCase):
     def test_company_environment(self):
         """ Check the company context on the environment is verified. """
 
-        self.company_c = self.env['res.company'].create({
-            'name': 'Company C'
-        })
-
-        user = self.env['res.users'].create({
-            'name': 'Test',
-            'login': 'test',
-            'company_id': self.company_a.id,
-            'company_ids': (self.company_a | self.company_c).ids,
-        })
-
-        user = user.with_user(user).with_context(allowed_company_ids=[])
+        user = self.test_user.with_user(self.test_user).with_context(allowed_company_ids=[])
 
         # When accessing company/companies, check raises error if unauthorized/unexisting company.
         with self.assertRaises(AccessError):
@@ -120,3 +119,57 @@ class TestCompanyCheck(common.TransactionCase):
         # Fallbacks when no allowed_company_ids context key
         self.assertEqual(user.env.company, user.company_id)
         self.assertEqual(user.env.companies, user.company_ids)
+
+    def test_with_company(self):
+        """ Check that with_company() works as expected """
+
+        user = self.test_user.with_user(self.test_user).with_context(allowed_company_ids=[])
+        self.assertEqual(user.env.companies, user.company_ids)
+
+        self.assertEqual(user.with_company(user.env.company).env.company, user.env.company)
+        self.assertEqual(
+            user.with_company(self.company_a).env.company,
+            user.with_company(self.company_a.id).env.company
+        )
+
+        # Falsy values shouldn't change current environment, record, or context.
+        for falsy in [False, None, 0, '', self.env['res.company'], []]:
+            no_change = user.with_company(falsy)
+            self.assertEqual(no_change, user)
+            self.assertEqual(no_change.env.context, user.env.context)
+            self.assertEqual(no_change.env, user.env)
+
+        comp_a_user = user.with_company(user.company_id)
+        comp_a_user2 = user.with_context(allowed_company_ids=user.company_id.ids)
+
+        # Using with_company(c) or with_context(allowed_company_ids=[c])
+        # should return the same context
+        # and the same environment (reused if no changes)
+        self.assertEqual(comp_a_user.env, comp_a_user2.env)
+        self.assertEqual(comp_a_user.env.context, comp_a_user2.env.context)
+
+        # When there were no company in the context, using with_company
+        # restricts both env.company and env.companies.
+        self.assertEqual(comp_a_user.env.company, user.company_id)
+        self.assertEqual(comp_a_user.env.companies, user.company_id)
+
+        # Reordering allowed_company_ids ctxt key
+        # Ensure with_company reorders the context key content
+        # and by consequent changes env.company
+        comp_c_a_user = comp_a_user.with_company(self.company_c)
+        self.assertEqual(comp_c_a_user.env.company, self.company_c)
+        self.assertEqual(comp_c_a_user.env.companies, self.company_c + self.company_a)
+        self.assertEqual(comp_c_a_user.env.companies[0], self.company_c)
+        self.assertEqual(comp_c_a_user.env.companies[1], self.company_a)
+        self.assertEqual(
+            comp_c_a_user.env.context['allowed_company_ids'],
+            [self.company_c.id, self.company_a.id],
+        )
+
+        comp_a_c_user = comp_c_a_user.with_company(self.company_a)
+        self.assertEqual(comp_a_c_user.env.companies[0], self.company_a)
+        self.assertEqual(comp_a_c_user.env.companies[1], self.company_c)
+        self.assertEqual(
+            comp_a_c_user.env.context['allowed_company_ids'],
+            [self.company_a.id, self.company_c.id],
+        )
diff --git a/odoo/models.py b/odoo/models.py
index 96f41eecc1d..b153b8c22a4 100644
--- a/odoo/models.py
+++ b/odoo/models.py
@@ -5011,6 +5011,39 @@ Record ids: %(records)s
         """
         return self.with_env(self.env(user=user, su=False))
 
+    def with_company(self, company):
+        """ with_company(company)
+
+        Return a new version of this recordset with a modified context, such that::
+            result.env.company = company
+            result.env.companies = self.env.companies | company
+
+        :param company:
+        :type company: :class:`~odoo.addons.base.models.res_company` or int
+
+        .. warning::
+
+            When using an unauthorized company for current user,
+            accessing the company(ies) on the environment may trigger
+            an AccessError if not done in a sudoed environment.
+        """
+        if not company:
+            # With company = None/False/0/[]/empty recordset: keep current environment
+            return self
+
+        company_id = int(company)
+        allowed_company_ids = self.env.context.get('allowed_company_ids', [])
+        if allowed_company_ids and company_id == allowed_company_ids[0]:
+            return self
+        # Copy the allowed_company_ids list
+        # to avoid modifying the context of the current environment.
+        allowed_company_ids = list(allowed_company_ids)
+        if company_id in allowed_company_ids:
+            allowed_company_ids.remove(company_id)
+        allowed_company_ids.insert(0, company_id)
+
+        return self.with_context(allowed_company_ids=allowed_company_ids)
+
     def with_context(self, *args, **kwargs):
         """ with_context([context][, **overrides]) -> records