chore: Remove __init__.py for batch materialization engines (feast-…

…dev#3133) * Remove `__init__.py` for batch materialization engines Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Do not execute `app.py` during doctests Signed-off-by: Felix Wang <wangfelix98@gmail.com> Signed-off-by: Felix Wang <wangfelix98@gmail.com>
srausser · Aug 24, 2022 · 5da89f5 · 5da89f5
1 parent 66ee5d8
commit 5da89f5
Show file tree

Hide file tree

Showing 10 changed files with 14 additions and 39 deletions.
diff --git a/sdk/python/docs/index.rst b/sdk/python/docs/index.rst
@@ -328,12 +328,12 @@ Cassandra Online Store
 Batch Materialization Engine
 ============================
 
-.. automodule:: feast.infra.materialization
+.. automodule:: feast.infra.materialization.batch_materialization_engine
     :members: BatchMaterializationEngine, MaterializationJob, MaterializationTask
 
 Local Engine
 ------------
-.. autoclass:: feast.infra.materialization.LocalMaterializationEngine
+.. autoclass:: feast.infra.materialization.local_engine.LocalMaterializationEngine
     :members:
     :noindex:
 

diff --git a/sdk/python/feast/infra/materialization/__init__.py b/sdk/python/feast/infra/materialization/__init__.py
diff --git a/sdk/python/feast/infra/materialization/aws_lambda/__init__.py b/sdk/python/feast/infra/materialization/aws_lambda/__init__.py
diff --git a/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py b/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py
@@ -13,7 +13,7 @@
 from feast import FeatureView, RepoConfig
 from feast.batch_feature_view import BatchFeatureView
 from feast.entity import Entity
-from feast.infra.materialization import (
+from feast.infra.materialization.batch_materialization_engine import (
     BatchMaterializationEngine,
     MaterializationJob,
     MaterializationTask,

diff --git a/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_task.py b/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_task.py
@@ -1,4 +1,4 @@
-from feast.infra.materialization import MaterializationTask
+from feast.infra.materialization.batch_materialization_engine import MaterializationTask
 
 
 class BytewaxMaterializationTask(MaterializationTask):

diff --git a/sdk/python/feast/infra/materialization/local_engine.py b/sdk/python/feast/infra/materialization/local_engine.py
@@ -12,12 +12,12 @@
 from feast.infra.registry.base_registry import BaseRegistry
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
 from feast.stream_feature_view import StreamFeatureView
-
-from ...utils import (
+from feast.utils import (
     _convert_arrow_to_proto,
     _get_column_names,
     _run_pyarrow_field_mapping,
 )
+
 from .batch_materialization_engine import (
     BatchMaterializationEngine,
     MaterializationJob,

diff --git a/sdk/python/feast/infra/passthrough_provider.py b/sdk/python/feast/infra/passthrough_provider.py
@@ -11,9 +11,10 @@
 from feast.feature_logging import FeatureServiceLoggingSource
 from feast.feature_service import FeatureService
 from feast.feature_view import FeatureView
-from feast.infra.materialization import BatchMaterializationEngine, MaterializationTask
 from feast.infra.materialization.batch_materialization_engine import (
+    BatchMaterializationEngine,
     MaterializationJobStatus,
+    MaterializationTask,
 )
 from feast.infra.offline_stores.offline_store import RetrievalJob
 from feast.infra.offline_stores.offline_utils import get_offline_store_from_config

diff --git a/sdk/python/feast/infra/transformation_servers/__init__.py b/sdk/python/feast/infra/transformation_servers/__init__.py
diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -35,7 +35,7 @@
 # - existing values for the online store type in featurestore.yaml files continue to work in a backwards compatible way
 # - first party and third party implementations can use the same class loading code path.
 BATCH_ENGINE_CLASS_FOR_TYPE = {
-    "local": "feast.infra.materialization.LocalMaterializationEngine",
+    "local": "feast.infra.materialization.local_engine.LocalMaterializationEngine",
     "snowflake.engine": "feast.infra.materialization.snowflake_engine.SnowflakeMaterializationEngine",
     "lambda": "feast.infra.materialization.aws_lambda.lambda_engine.LambdaMaterializationEngine",
     "bytewax": "feast.infra.materialization.contrib.bytewax.bytewax_materialization_engine.BytewaxMaterializationEngine",

diff --git a/sdk/python/tests/doctest/test_all.py b/sdk/python/tests/doctest/test_all.py
@@ -7,6 +7,8 @@
 
 import feast
 
+FILES_TO_IGNORE = {"app"}
+
 
 def setup_feature_store():
     """Prepares the local environment for a FeatureStore docstring test."""
@@ -69,6 +71,9 @@ def test_docstrings():
 
         for package in current_packages:
             for _, name, is_pkg in pkgutil.walk_packages(package.__path__):
+                if name in FILES_TO_IGNORE:
+                    continue
+
                 full_name = package.__name__ + "." + name
                 try:
                     temp_module = importlib.import_module(full_name)