[PYG-263, PYG-277] 🧞Connection shortcut (#348)

* feat: added option for dropping na columns

* refactor: fix correct template

* refactor: regen

* refactor: setup connection Item A example

* refactor: new template

* refactor; regen

* fix: handle edeg class

* fix: skip non writeable linked class

* refactor; regen

* fix: allow reverse list

* fix: allow reverse of list

* refactor: added core model to examples

* refactor: download core

* feat: multiproces SDK gen

* refactor; gen CogniteCore

* refactor: increase to 8 processors

* refactor: revert qurey property

* build: changelog

* tests; query property to method

cognite/pygen/_core/models/data_classes.py

@@ -332,6 +332,25 @@ def write_connection_fields(self) -> Iterator[BaseConnectionField]:
             if isinstance(field, BaseConnectionField) and not field.is_direct_relation_no_source
         )
 
+    @property
+    def write_connection_fields_including_parents(self) -> Iterator[BaseConnectionField]:
+        """These fields are used when creating properties in the lists classes"""
+        return (
+            field
+            for field in self
+            if isinstance(field, BaseConnectionField)
+            and not field.is_direct_relation_no_source
+            and field.is_write_field
+            and field.linked_class.is_writable
+        )
+
+    @property
+    def read_connection_fields_including_parents(self) -> Iterator[BaseConnectionField]:
+        """These fields are used when creating properties in the lists classes as well as querying"""
+        return (
+            field for field in self if isinstance(field, BaseConnectionField) and not field.is_direct_relation_no_source
+        )
+
     @property
     def has_write_connection_fields(self) -> bool:
         """Check if the data class has any write connection fields."""

cognite/pygen/_core/templates/api_class_node.py.jinja

@@ -408,7 +408,6 @@ class {{ api_class.name }}({% if data_class.is_writable %}NodeAPI{% else %}NodeR
             filter_,
         )
 
-    @property
     def query(self) -> {{ data_class.query_cls_name }}:
         """Start a query for {{ data_class.doc_list_name}}."""
         warnings.warn("The .query is in alpha and is subject to breaking changes without notice.")

cognite/pygen/_core/templates/data_class_node.py.jinja

@@ -56,7 +56,7 @@ from {{ top_level_package }}.data_classes._core import ({% if has_default_instan
 from {{ top_level_package }}.data_classes.{{ implements.file_name }} import {{ implements.read_name }}{% if data_class.is_writable or data_class.is_interface %}, {{ implements.write_name }}{% endif %}{% endfor %}{% endif %}
 {% if data_class.has_dependencies_not_self %}
 if TYPE_CHECKING:{% for dependency_class in data_class.dependencies %}{% if dependency_class.file_name != data_class.file_name %}
-    from {{ top_level_package }}.data_classes.{{ dependency_class.file_name }} import {{ dependency_class.read_name }}, {{ dependency_class.graphql_name }}{% if dependency_class.is_writable or dependency_class.is_interface %}, {{ dependency_class.write_name }}{% endif %}{% endif %}{% endfor %}
+    from {{ top_level_package }}.data_classes.{{ dependency_class.file_name }} import {{ dependency_class.read_name }}, {{ dependency_class.read_list_name }}, {{ dependency_class.graphql_name }}{% if dependency_class.is_writable or dependency_class.is_interface %}, {{ dependency_class.write_name }}, {{ dependency_class.write_list_name }}{% endif %}{% endif %}{% endfor %}
 {% endif %}
 
 __all__ = [
@@ -487,13 +487,27 @@ class {{ data_class.read_list_name }}(DomainModelList[{{ data_class.read_name }}
             stacklevel=2,
         )
         return self.as_write()
-
+{% for field in data_class.read_connection_fields_including_parents %}
+    @property
+    def {{ field.name }}(self) -> {{ field.linked_class.read_list_name }}:{% if field.linked_class.file_name != data_class.file_name %}
+        from .{{ field.linked_class.file_name }} import {{ field.linked_class.read_name }}, {{ field.linked_class.read_list_name }}
+{% endif %}{% if field.is_one_to_one %}
+        return {{ field.linked_class.read_list_name }}([item.{{ field.name }} for item in self.data if isinstance(item.{{ field.name }}, {{ field.linked_class.read_name }})]){% else %}
+        return {{ field.linked_class.read_list_name }}([item for items in self.data for item in items.{{ field.name }} or [] if isinstance(item, {{ field.linked_class.read_name }})]){% endif %}
+{% endfor %}
 
 class {{ data_class.write_list_name }}(DomainModelWriteList[{{ data_class.write_name }}]):
     """List of {{ data_class.doc_list_name }} in the writing version."""
 
     _INSTANCE = {{ data_class.write_name }}
-
+{% for field in data_class.write_connection_fields_including_parents %}
+    @property
+    def {{ field.name }}(self) -> {{ field.linked_class.write_list_name }}:{% if field.linked_class.file_name != data_class.file_name %}
+        from .{{ field.linked_class.file_name }} import {{ field.linked_class.write_name }}, {{ field.linked_class.write_list_name }}
+{% endif %}{% if field.is_one_to_one %}
+        return {{ field.linked_class.write_list_name }}([item.{{ field.name }} for item in self.data if isinstance(item.{{ field.name }}, {{ field.linked_class.write_name }})]){% else %}
+        return {{ field.linked_class.write_list_name }}([item for items in self.data for item in items.{{ field.name }} or [] if isinstance(item, {{ field.linked_class.write_name }})]){% endif %}
+{% endfor %}
 class {{ data_class.read_name }}ApplyList({{ data_class.write_list_name }}): ...
 
 {% endif %}
@@ -541,7 +555,7 @@ class _{{ data_class.query_cls_name }}(NodeQueryCore[T_DomainModelList, {{ data_
             reverse_expression,
         )
 {% for field in data_class.fields_of_type(ft.BaseConnectionField) %}{% if not field.is_direct_relation_no_source %}
-        if _{{ field.linked_class.query_cls_name }} not in created_types and connection_type != "reverse-list":
+        if _{{ field.linked_class.query_cls_name }} not in created_types:
             self.{{ field.name }} = _{{ field.linked_class.query_cls_name }}(
                 created_types.copy(),
                 self._creation_path,

cognite/pygen/_core/templates/data_classes_core_base.py.jinja

@@ -144,10 +144,13 @@ class GraphQLList(UserList):
     def dump(self) -> list[dict[str, Any]]:
         return [node.model_dump() for node in self.data]
 
-    def to_pandas(self) -> pd.DataFrame:
+    def to_pandas(self, dropna_columns: bool = False) -> pd.DataFrame:
         """
         Convert the list of nodes to a pandas.DataFrame.
 
+        Args:
+            dropna_columns: Whether to drop columns that are all NaN.
+
         Returns:
             A pandas.DataFrame with the nodes as rows.
         """
@@ -161,10 +164,15 @@ class GraphQLList(UserList):
         columns = (
             id_columns + [col for col in df if col not in fixed_columns] + [col for col in end_columns if col in df]
         )
-        return df[columns]
+        df = df[columns]
+        if df.empty:
+            return df
+        if dropna_columns:
+            df.dropna(how="all", axis=1, inplace=True)
+        return df
 
     def _repr_html_(self) -> str:
-        return self.to_pandas()._repr_html_()  # type: ignore[operator]
+        return self.to_pandas(dropna_columns=True)._repr_html_()  # type: ignore[operator]
 
 
 class DomainModelCore(Core, ABC):
@@ -397,10 +405,13 @@ class CoreList(UserList, Generic[T_Core]):
     def as_external_ids(self) -> list[str]:
         return [node.external_id for node in self.data]
 
-    def to_pandas(self) -> pd.DataFrame:
+    def to_pandas(self, dropna_columns: bool = False) -> pd.DataFrame:
         """
         Convert the list of nodes to a pandas.DataFrame.
 
+        Args:
+            dropna_columns: Whether to drop columns that are all NaN.
+
         Returns:
             A pandas.DataFrame with the nodes as rows.
         """
@@ -414,10 +425,15 @@ class CoreList(UserList, Generic[T_Core]):
         columns = (
             id_columns + [col for col in df if col not in fixed_columns] + [col for col in end_columns if col in df]
         )
-        return df[columns]
+        df = df[columns]
+        if df.empty:
+            return df
+        if dropna_columns:
+            df.dropna(how="all", axis=1, inplace=True)
+        return df
 
     def _repr_html_(self) -> str:
-        return self.to_pandas()._repr_html_()  # type: ignore[operator]
+        return self.to_pandas(dropna_columns=True)._repr_html_()  # type: ignore[operator]
 
 
 class DomainModelList(CoreList[T_DomainModel]):

dev.py

@@ -3,6 +3,7 @@
 import re
 from collections import defaultdict
 from datetime import datetime
+from multiprocessing import Pool
 from typing import TypeVar
 
 import toml
@@ -14,7 +15,7 @@
 from cognite.pygen._generator import SDKGenerator, generate_typed, write_sdk_to_disk
 from cognite.pygen.utils import MockGenerator
 from cognite.pygen.utils.cdf import load_cognite_client_from_toml
-from tests.constants import DATA_WRITE_DIR, EXAMPLE_SDKS, EXAMPLES_DIR, REPO_ROOT
+from tests.constants import DATA_WRITE_DIR, EXAMPLE_SDKS, EXAMPLES_DIR, REPO_ROOT, ExampleSDK
 
 app = typer.Typer(
     add_completion=False,
@@ -29,45 +30,53 @@
 def generate_sdks(
     sdk_name: str = typer.Option(None, "--sdk", help="Generate only the specified SDK"),
 ):
-    for example_sdk in EXAMPLE_SDKS:
-        if not example_sdk.generate_sdk:
-            continue
-        if sdk_name is not None and not example_sdk.client_name.casefold().startswith(sdk_name.casefold()):
-            continue
-        typer.echo(f"Generating {example_sdk.client_name} SDK...")
-        data_models = example_sdk.load_data_models()
-        if len(data_models) == 1:
-            data_models = data_models[0]
-
-        if example_sdk.is_typed:
-            output_file = example_sdk.client_dir / "typed.py"
-            include_views = {dm.ViewId(data_models.space, t, "1") for t in example_sdk.typed_classes} or None
-            generate_typed(data_models, output_file, include_views=include_views, implements="composition")
-            continue
-
-        sdk_generator = SDKGenerator(
-            example_sdk.top_level_package,
-            example_sdk.client_name,
-            data_models,
-            logger=typer.echo,
-            default_instance_space=example_sdk.instance_space,
+    sdks_to_generate = (example_sdk for example_sdk in EXAMPLE_SDKS if example_sdk.generate_sdk)
+    if sdk_name is not None:
+        sdks_to_generate = (
+            example_sdk
+            for example_sdk in sdks_to_generate
+            if not example_sdk.client_name.casefold().startswith(sdk_name.casefold())
         )
-
-        sdk = sdk_generator.generate_sdk()
-        write_sdk_to_disk(
-            sdk,
-            example_sdk.client_dir,
-            overwrite=True,
-            logger=print,
-            format_code=True,
-        )
-        typer.echo(f"{example_sdk.client_name} SDK Created in {example_sdk.client_dir}")
-        typer.echo("All files updated! Including files assumed to be manually maintained.")
-        typer.echo("\n")
+    sdks = list(sdks_to_generate)
+    with Pool(min(8, len(sdks))) as pool:
+        pool.map(_generate_sdk, sdks)
 
     typer.echo("All SDKs Created!")
 
 
+def _generate_sdk(example_sdk: ExampleSDK) -> None:
+    typer.echo(f"Generating {example_sdk.client_name} SDK...")
+    data_models = example_sdk.load_data_models()
+    if len(data_models) == 1:
+        data_models = data_models[0]
+
+    if example_sdk.is_typed:
+        output_file = example_sdk.client_dir / "typed.py"
+        include_views = {dm.ViewId(data_models.space, t, "1") for t in example_sdk.typed_classes} or None
+        generate_typed(data_models, output_file, include_views=include_views, implements="composition")
+        return
+
+    sdk_generator = SDKGenerator(
+        example_sdk.top_level_package,
+        example_sdk.client_name,
+        data_models,
+        logger=typer.echo,
+        default_instance_space=example_sdk.instance_space,
+    )
+
+    sdk = sdk_generator.generate_sdk()
+    write_sdk_to_disk(
+        sdk,
+        example_sdk.client_dir,
+        overwrite=True,
+        logger=print,
+        format_code=True,
+    )
+    typer.echo(f"{example_sdk.client_name} SDK Created in {example_sdk.client_dir}")
+    typer.echo("All files updated! Including files assumed to be manually maintained.")
+    typer.echo("\n")
+
+
 @app.command("download", help="Download the DMS representation of all example SDKs")
 def download():
     client = load_cognite_client_from_toml("config.toml")

docs/CHANGELOG.md

@@ -18,6 +18,15 @@ Changes are grouped as follows
 - When calling `generate_sdk` or `generate_sdk_notebook` you now get an improved formatting of warnings. In addition,
   irrelevant warnings are no longer shown.
 
+### Added
+- In the generated list classes, all connection properties are now available as properties. This is to make it easier
+  to access the connection properties without having to go through the `data_record` property.
+
+### Changed
+- [Experimental - Breaking] Reverted breaking change from `0.99.40`; The `.query` property is now a method `query`.
+  Instead of `client.asset.query` you now do `client.asset.query()`. This is to make type hints better when working
+  ina Notebook environment.
+
 ## [0.99.44] - 24-11-08
 ### Improved
 - In the underlying `query` operation, used for `.list(...)` as well as query methods, `pygen` now

examples/cognite_core/__init__.py

@@ -0,0 +1,3 @@
+from cognite_core._api_client import CogniteCoreClient
+
+__all__ = ["CogniteCoreClient"]

examples/cognite_core/_api/__init__.py

@@ -0,0 +1,125 @@
+from cognite_core._api.cognite_360_image import Cognite360ImageAPI
+from cognite_core._api.cognite_360_image_collection import Cognite360ImageCollectionAPI
+from cognite_core._api.cognite_360_image_collection_query import Cognite360ImageCollectionQueryAPI
+from cognite_core._api.cognite_360_image_model import Cognite360ImageModelAPI
+from cognite_core._api.cognite_360_image_model_query import Cognite360ImageModelQueryAPI
+from cognite_core._api.cognite_360_image_query import Cognite360ImageQueryAPI
+from cognite_core._api.cognite_360_image_station import Cognite360ImageStationAPI
+from cognite_core._api.cognite_360_image_station_query import Cognite360ImageStationQueryAPI
+from cognite_core._api.cognite_3_d_model import Cognite3DModelAPI
+from cognite_core._api.cognite_3_d_model_query import Cognite3DModelQueryAPI
+from cognite_core._api.cognite_3_d_object import Cognite3DObjectAPI
+from cognite_core._api.cognite_3_d_object_images_360 import Cognite3DObjectImagesAPI
+from cognite_core._api.cognite_3_d_object_query import Cognite3DObjectQueryAPI
+from cognite_core._api.cognite_3_d_revision import Cognite3DRevisionAPI
+from cognite_core._api.cognite_3_d_revision_query import Cognite3DRevisionQueryAPI
+from cognite_core._api.cognite_3_d_transformation_node import Cognite3DTransformationNodeAPI
+from cognite_core._api.cognite_3_d_transformation_node_query import Cognite3DTransformationNodeQueryAPI
+from cognite_core._api.cognite_activity import CogniteActivityAPI
+from cognite_core._api.cognite_activity_query import CogniteActivityQueryAPI
+from cognite_core._api.cognite_asset import CogniteAssetAPI
+from cognite_core._api.cognite_asset_class import CogniteAssetClassAPI
+from cognite_core._api.cognite_asset_class_query import CogniteAssetClassQueryAPI
+from cognite_core._api.cognite_asset_query import CogniteAssetQueryAPI
+from cognite_core._api.cognite_asset_type import CogniteAssetTypeAPI
+from cognite_core._api.cognite_asset_type_query import CogniteAssetTypeQueryAPI
+from cognite_core._api.cognite_cad_model import CogniteCADModelAPI
+from cognite_core._api.cognite_cad_model_query import CogniteCADModelQueryAPI
+from cognite_core._api.cognite_cad_node import CogniteCADNodeAPI
+from cognite_core._api.cognite_cad_node_query import CogniteCADNodeQueryAPI
+from cognite_core._api.cognite_cad_revision import CogniteCADRevisionAPI
+from cognite_core._api.cognite_cad_revision_query import CogniteCADRevisionQueryAPI
+from cognite_core._api.cognite_cube_map import CogniteCubeMapAPI
+from cognite_core._api.cognite_cube_map_query import CogniteCubeMapQueryAPI
+from cognite_core._api.cognite_describable_node import CogniteDescribableNodeAPI
+from cognite_core._api.cognite_describable_node_query import CogniteDescribableNodeQueryAPI
+from cognite_core._api.cognite_equipment import CogniteEquipmentAPI
+from cognite_core._api.cognite_equipment_query import CogniteEquipmentQueryAPI
+from cognite_core._api.cognite_equipment_type import CogniteEquipmentTypeAPI
+from cognite_core._api.cognite_equipment_type_query import CogniteEquipmentTypeQueryAPI
+from cognite_core._api.cognite_file import CogniteFileAPI
+from cognite_core._api.cognite_file_category import CogniteFileCategoryAPI
+from cognite_core._api.cognite_file_category_query import CogniteFileCategoryQueryAPI
+from cognite_core._api.cognite_file_query import CogniteFileQueryAPI
+from cognite_core._api.cognite_point_cloud_model import CognitePointCloudModelAPI
+from cognite_core._api.cognite_point_cloud_model_query import CognitePointCloudModelQueryAPI
+from cognite_core._api.cognite_point_cloud_revision import CognitePointCloudRevisionAPI
+from cognite_core._api.cognite_point_cloud_revision_query import CognitePointCloudRevisionQueryAPI
+from cognite_core._api.cognite_point_cloud_volume import CognitePointCloudVolumeAPI
+from cognite_core._api.cognite_point_cloud_volume_query import CognitePointCloudVolumeQueryAPI
+from cognite_core._api.cognite_schedulable import CogniteSchedulableAPI
+from cognite_core._api.cognite_schedulable_query import CogniteSchedulableQueryAPI
+from cognite_core._api.cognite_source_system import CogniteSourceSystemAPI
+from cognite_core._api.cognite_source_system_query import CogniteSourceSystemQueryAPI
+from cognite_core._api.cognite_sourceable_node import CogniteSourceableNodeAPI
+from cognite_core._api.cognite_sourceable_node_query import CogniteSourceableNodeQueryAPI
+from cognite_core._api.cognite_time_series import CogniteTimeSeriesAPI
+from cognite_core._api.cognite_time_series_query import CogniteTimeSeriesQueryAPI
+from cognite_core._api.cognite_unit import CogniteUnitAPI
+from cognite_core._api.cognite_unit_query import CogniteUnitQueryAPI
+from cognite_core._api.cognite_visualizable import CogniteVisualizableAPI
+from cognite_core._api.cognite_visualizable_query import CogniteVisualizableQueryAPI
+
+__all__ = [
+    "Cognite360ImageAPI",
+    "Cognite360ImageCollectionAPI",
+    "Cognite360ImageCollectionQueryAPI",
+    "Cognite360ImageModelAPI",
+    "Cognite360ImageModelQueryAPI",
+    "Cognite360ImageQueryAPI",
+    "Cognite360ImageStationAPI",
+    "Cognite360ImageStationQueryAPI",
+    "Cognite3DModelAPI",
+    "Cognite3DModelQueryAPI",
+    "Cognite3DObjectAPI",
+    "Cognite3DObjectImagesAPI",
+    "Cognite3DObjectQueryAPI",
+    "Cognite3DRevisionAPI",
+    "Cognite3DRevisionQueryAPI",
+    "Cognite3DTransformationNodeAPI",
+    "Cognite3DTransformationNodeQueryAPI",
+    "CogniteActivityAPI",
+    "CogniteActivityQueryAPI",
+    "CogniteAssetAPI",
+    "CogniteAssetClassAPI",
+    "CogniteAssetClassQueryAPI",
+    "CogniteAssetQueryAPI",
+    "CogniteAssetTypeAPI",
+    "CogniteAssetTypeQueryAPI",
+    "CogniteCADModelAPI",
+    "CogniteCADModelQueryAPI",
+    "CogniteCADNodeAPI",
+    "CogniteCADNodeQueryAPI",
+    "CogniteCADRevisionAPI",
+    "CogniteCADRevisionQueryAPI",
+    "CogniteCubeMapAPI",
+    "CogniteCubeMapQueryAPI",
+    "CogniteDescribableNodeAPI",
+    "CogniteDescribableNodeQueryAPI",
+    "CogniteEquipmentAPI",
+    "CogniteEquipmentQueryAPI",
+    "CogniteEquipmentTypeAPI",
+    "CogniteEquipmentTypeQueryAPI",
+    "CogniteFileAPI",
+    "CogniteFileCategoryAPI",
+    "CogniteFileCategoryQueryAPI",
+    "CogniteFileQueryAPI",
+    "CognitePointCloudModelAPI",
+    "CognitePointCloudModelQueryAPI",
+    "CognitePointCloudRevisionAPI",
+    "CognitePointCloudRevisionQueryAPI",
+    "CognitePointCloudVolumeAPI",
+    "CognitePointCloudVolumeQueryAPI",
+    "CogniteSchedulableAPI",
+    "CogniteSchedulableQueryAPI",
+    "CogniteSourceSystemAPI",
+    "CogniteSourceSystemQueryAPI",
+    "CogniteSourceableNodeAPI",
+    "CogniteSourceableNodeQueryAPI",
+    "CogniteTimeSeriesAPI",
+    "CogniteTimeSeriesQueryAPI",
+    "CogniteUnitAPI",
+    "CogniteUnitQueryAPI",
+    "CogniteVisualizableAPI",
+    "CogniteVisualizableQueryAPI",
+]