feat: faker contacts

service/faker_data.py

@@ -1,3 +1,4 @@
+import copy
 import random
 import requests
 import threading
@@ -5,6 +6,8 @@ import random
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from lxml import etree
 from faker import Faker
+
+import utils.model
 from model import SexEnum, Person, Hotel, Railway, City
 from base import db
 from datetime import datetime, timedelta
@@ -78,7 +81,9 @@ def faker_identifications(num: int) -> list[Person]:
     persons = [faker_identification() for i in range(num)]
     area_codes = [person.identification[0:6] for person in persons]
     rows = City.query.filter(City.ad_code.in_(area_codes))
-    code_area_dict = {row.ad_code: '{} {} {} {}'.format(row.country_cn, row.province_cn, row.admin_district_cn, row.city_cn) for row in rows}
+    code_area_dict = {
+        row.ad_code: '{} {} {} {}'.format(row.country_cn, row.province_cn, row.admin_district_cn, row.city_cn) for row
+        in rows}
     for person in persons:
         person.address = code_area_dict[person.identification[0:6]]
     db.session.bulk_save_objects(persons)
@@ -108,4 +113,61 @@ def generate_train_id():
     num = int(random.random() * 10000).__str__()
     return prefix + num
 
-# def faker_hotel(num: int):
+
+def faker_hotels(num: int) -> list[Hotel]:
+    persons = faker_identifications(num)
+    hotels = [faker_hotel(person) for person in persons]
+    db.session.bulk_save_objects(hotels)
+    db.session.commit()
+    return hotels
+
+
+def faker_hotel(person: Person) -> Hotel:
+    hotel = faker_hotel_infos[random.randint(0, len(faker_hotel_infos) - 1)]
+    hotel.identification = person.identification
+    hotel.in_data = datetime.today() + timedelta(days=(int(random.random() * 21) - 10))
+    hotel.out_data = hotel.in_data + timedelta(days=3)
+    return hotel
+
+
+def faker_contacts(num_hotel: int, num_railway: int):
+    persons = faker_identifications(num_railway + num_hotel + 1)
+    patient = persons[0]
+    contacts_hotel = persons[1:1 + num_hotel]
+    contacts_railway = persons[-1-num_railway: -1]
+    print(persons)
+    print(patient)
+    print(contacts_railway)
+    print(contacts_hotel)
+    hotel_patient = faker_hotel(patient)
+    railway_patient = faker_railway(patient)
+
+    def assemble_hotel(contact: Person) -> Hotel:
+        hotel_contact = copy.copy(hotel_patient)
+        hotel_contact.identification = contact.identification
+        return hotel_contact
+
+    def assemble_railway(contact: Person) -> Railway:
+        railway_contact = copy.copy(railway_patient)
+        railway_contact.identification = contact.identification
+        railway_contact.name = contact.name
+        return railway_contact
+
+    hotel_contacts = [assemble_hotel(contact) for contact in contacts_hotel]
+    railway_contacts = [assemble_railway(contact) for contact in contacts_railway]
+    hotels = [hotel_patient] + hotel_contacts
+    railways = [railway_patient] + railway_contacts
+    db.session.bulk_save_objects(hotels)
+    db.session.bulk_save_objects(railways)
+    db.session.commit()
+    return {
+        "patient": {
+            "hotel": utils.model.model2dict(hotel_patient),
+            "railway": utils.model.model2dict(railway_patient)
+        },
+        "contacts": {
+            "hotel": utils.model.models2dicts(hotel_contacts),
+            "railway": utils.model.models2dicts(railway_contacts)
+        }
+    }
+