How to write a peewee class that successfully initialises other peewee objects during its own initialisation?

Question

The classes I have are:

class B(Base_Model):
    b_attribute_1 = peewee.TextField(null=True)
    a = peewee.ForeignKey(A, null=True)

class A(Base_Model):
    a_attribute_1 = peewee.IntegerField(null=True)
    a_attribute_2 = peewee.DoubleField(null=True)
    class Meta:
        friend_server_address = "103.11.399.002"
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.save()
        from B import B
        # Option (1) -- semi-works: creates incomplete object b,
        #               a is initialised correctly.
        B.create(a=self)
        # Option (2) -- does not work: b is created successfully,
        #               a not created, giving KeyError = "ba1".
        # B.create(b_attribute_1 = kwargs["ba1"], a=self)
        return

The way I see problems described in code comments (see Options 1 and 2) is by testing the code:

class test_a_and_b(unittest.TestCase):
    def setUp(self):
        config.database.init("test.db")
        config.database.create_tables([A, B])
        A(a_attribute_1 = 9705,
          a_attribute_2 = 0.77,
          ba1 = "this is b")
        # This is the breakpoint where I study objects in memory.
        pdb.set_trace()
        return
    def tearDownModule()
        A.get().delete_instance(recursive=True, delete_nullable=True)  
        B.get().delete_instance(recursive=True, delete_nullable=True)
        config.database.drop_tables([A, B])
        os.remove("test.db")

So depending on the option I choose in __init__ of my A I (after calling A.get() and B.get() at break point) can see that either complete a and incomplete b are created, or just complete b and error message KeyError:"ba1". What I want is two complete objects (database rows) a and b being created.

I tried looking into the peewee library, it seems their Model uses a __new__, I did not see any __init__, so perhaps here lies my problem. Tried different combos in my A, but that did not seem to help. Also tried using initialise, define attributes programmatically and .save() instead .create() for b, but that did not help either.

Answer 1

Pretty sure the problem that you're having is that you're not always going to be including ba1 when you want an instance of A. It's only when you're creating a new one that you'll be passing in ba1. When you select one or when peewee internally instantiates one when referencing it from B, your 'ba1' keyword argument won't be there and your code will fail with the error that you're getting.

Check for the existence of 'ba1' before using it (as you will only be using it when you are explicitly creating a new A yourself). In other words, try this:

import peewee
import unittest

database_proxy = peewee.Proxy()
class BaseModel(peewee.Model):
    class Meta:
        database = database_proxy

class A(BaseModel):
    a_attribute_1 = peewee.IntegerField(null=True)
    a_attribute_2 = peewee.IntegerField(null=True)
    def __init__(self, *args, **kwargs):
        super(A, self).__init__(*args, **kwargs)
        if "ba1" in kwargs:
            self.save()
            B.create(a=self, b_attribute_1 = kwargs["ba1"])

class B(BaseModel):
    b_attribute_1 = peewee.TextField(null=True)
    a = peewee.ForeignKeyField(A, null=True)

class test_a_and_b(unittest.TestCase):
    def setUp(self):
        db = peewee.SqliteDatabase(':memory:')
        database_proxy.initialize(db)
        db.create_tables([A, B])
        A(a_attribute_1 = 9705,
          a_attribute_2 = 0.77,
          ba1 = "this is b")

    def test_a(self):
        b = B.select().where(B.b_attribute_1 == 'this is b').get()
        self.assertEqual(b.b_attribute_1, 'this is b')
        self.assertEqual(b.a.a_attribute_1, 9705)


if __name__ == '__main__':
    unittest.main()