Coverage for core / src / sensorkit / webapi / schema.py: 100%

1# SPDX-License-Identifier: Apache-2.0 

2import itertools 

3from typing import Any 

4 

5from pydantic import BaseModel, TypeAdapter 

6 

7from sensorkit import api as sk 

8from sensorkit.common.keyword import _keyword_index 

9 

10 

11def _sensorkit_models(): 

12 # Return all models registered in the sensorkit package. 

13 # FIXME: This is top-level and should probably live in `sensorkit.api` 

14 return itertools.chain( 

15 _keyword_index, 

16 sk.DeviceCommand.registry.entries, 

17 sk.Task.registry.entries, 

18 ) 

19 

20 

21def _update_refs(obj): 

22 """Recursively update $ref from #/$defs/Name to #/components/schemas/Name.""" 

23 if isinstance(obj, dict): 

24 if "$ref" in obj and isinstance(obj["$ref"], str): 

25 if obj["$ref"].startswith("#/$defs/"): 

26 obj["$ref"] = obj["$ref"].replace("#/$defs/", "#/components/schemas/") 

27 for v in obj.values(): 

28 _update_refs(v) 

29 elif isinstance(obj, list): 

30 for item in obj: 

31 _update_refs(item) 

32 

33 

34def _add_model_schema(model: type[BaseModel], schemas: dict[str, Any]): 

35 adapter = TypeAdapter(model) 

36 schema = adapter.json_schema() 

37 

38 # Extract internal definitions and move them to global schemas. 

39 if "$defs" in schema: 

40 defs = schema.pop("$defs") 

41 

42 for name, def_schema in defs.items(): 

43 if name not in schemas: 

44 # Recursively update refs in the definition itself. 

45 _update_refs(def_schema) 

46 schemas[name] = def_schema 

47 

48 # Update refs in the main schema. 

49 _update_refs(schema) 

50 

51 # Add the class schema to components/schemas. 

52 schema_name = model.__name__ 

53 

54 if schema_name not in schemas: 

55 schemas[schema_name] = schema 

56 

57 

58def add_sensorkit_schema(schemas: dict[str, Any]): 

59 # Generate JSON schema for registered models. 

60 for cls in _sensorkit_models(): 

61 if issubclass(cls, BaseModel): 

62 _add_model_schema(cls, schemas)