framework/stability_process_predicate/usecases/stability_check_usecase.py

import numpy as np
import pybullet as p
import time
import pybullet_data
import os
import json


class StabilityCheckUseCase:
    def call(self, outPath: str, buildNumber: int, duration=500):
        DURATION = duration
        try:
            # Получение сгенерированного URDF для сборки
            assemblyUrdf = json.loads(
                (open(outPath + 'urdf-generation.json')).read()).get(buildNumber)
        except:
            return TypeError('not found urfd file or not found build number')
        inc = 0
        urdfs = []

        for el in assemblyUrdf:
            inc += 1
            file_to_open = outPath + str(inc) + '.urdf'

            f = open(file_to_open, 'w', encoding='utf-8',
                     errors='ignore')
            f.write(el)
            urdfs.append(os.path.abspath(f.name))
            f.close()

        p.connect(p.DIRECT)

        p.setGravity(0, 0, -10)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        bulletIds = []
        for el in urdfs:
            bulletIds.append(p.loadURDF(el))
        p.loadURDF("plane.urdf")
        resultCoords = []
        for i in range(DURATION):
            if (i + 200 == DURATION):
                inc = 0
                for el in bulletIds:
                    inc += 1
                    # Получение расположения деталей
                    pos, rot = p.getBasePositionAndOrientation(el)
                    resultCoords.append({
                        'id': inc,
                        "quaternion": rot,
                        "position": pos
                    })
            p.stepSimulation()
            time.sleep(1./240.)

        file_to_open = outPath + buildNumber + "_" + 'stability_coords.json'

        f = open(file_to_open, 'w', encoding='utf-8',errors='ignore')
        # Запись результатов тестирования
        f.write(json.dumps(resultCoords))
        for el in urdfs:
            os.remove(el)
        f.close()
        p.disconnect()