Adding company field to inventory

inventory.py

@@ -26,6 +26,10 @@ class Inventory(OSV):
     moves = fields.Many2Many(
         'stock.move', 'inventory_move_rel', 'inventory', 'move',
         'Generated moves')
+    company = fields.Many2One(
+        'company.company', 'Company', required=True, states=STATES,)
+
+
     state = fields.Selection(
         [('open','Open'),
          ('done','Done'),
@@ -47,6 +51,13 @@ class Inventory(OSV):
     def default_date(self, cursor, user, context=None):
         return datetime.datetime.today()
 
+    def default_company(self, cursor, user, context=None):
+        if context.get('company'):
+            return context.get('company')
+        user_obj = self.pool.get('res.user')
+        user = user_obj.browse(cursor, user, user, context=context)
+        return user.company.id
+
     def set_state_cancel(self, cursor, user, ids, context=None):
         self.write(cursor, user, ids, {
             'state': 'cancel',
@@ -57,7 +68,6 @@ class Inventory(OSV):
             'state': 'open',
             }, context=context)
 
-
     def set_state_done(self, cursor, user, ids, context=None):
         self.write(cursor, user, ids, {
             'state': 'done',
@@ -226,6 +236,7 @@ class CompleteInventory(Wizard):
                     {'product': product,
                      'uom': uom,
                      'quantity': qty,
+                     'company': inventory.company.id,
                      'inventory': inventory.id,},
                     context=context)
         return {}

inventory.xml

@@ -11,21 +11,23 @@
             <field name="type">form</field>
             <field name="arch" type="xml">
                 <![CDATA[
-                    <form string="Inventory" col="6">
+                    <form string="Inventory" col="4">
                         <label name="location"/>
                         <field name="location"/>
                         <label name="lost_found"/>
                         <field name="lost_found"/>
                         <label name="date"/>
                         <field name="date"/>
+                        <label name="company"/>
+                        <field name="company"/>
                         <label colspan="4"/>
                         <button string="Complete Inventory" type="action"
                             name="%(wizard_complete_inventory)d"
                             states="{'readonly': '''state != 'open' '''}"
-                            colspan="2"
+                            colspan="1"
                             help="Add an inventory line for each missing products"/>
-                        <field name="lines" colspan="6"/>
-                        <group col="4" colspan="6">
+                        <field name="lines" colspan="4"/>
+                        <group col="4" colspan="4">
                             <label name="state"/>
                             <field name="state"/>
                             <group colspan="2" col="3">
@@ -54,6 +56,7 @@
                         <field name="location" select="1"/>
                         <field name="date" select="1"/>
                         <field name="state" select="1"/>
+                        <field name="company" select="2"/>
                     </tree>
                 ]]>
             </field>