Better CLI loading for ServiceInterface implementations

src/ClassAttributes/ShouldLoadInCliContext.php

@@ -0,0 +1,21 @@
+<?php
+
+/**
+ * An attribute class used to indicate a ServiceInterface should be loaded in WP-CLI.
+ *
+ * @package EightshiftLibs\ClassAttributes
+ */
+
+namespace EightshiftLibs\ClassAttributes;
+
+use Attribute;
+
+/**
+ * A class attribute definition class, inspected with Reflection when setting up services for DI in Main.
+ * ServiceInterface classes annotated with this attribute should be loaded in the CLI context as well,
+ * although they don't have to implement ServiceCliInterface themselves.
+ */
+#[Attribute]
+class ShouldLoadInCliContext
+{
+}

src/Main/AbstractMain.php

@@ -14,10 +14,12 @@
 use DI\ContainerBuilder;
 use DI\Definition\Helper\AutowireDefinitionHelper;
 use DI\Definition\Reference;
+use EightshiftLibs\ClassAttributes\ShouldLoadInCliContext;
 use EightshiftLibs\Services\ServiceInterface;
 use EightshiftLibs\Services\ServiceCliInterface;
 // phpcs:ignore SlevomatCodingStandard.Namespaces.UnusedUses.UnusedUse
 use Exception;
+use ReflectionClass;
 
 /**
  * The main start class.
@@ -71,14 +73,29 @@ public function registerServices()
 		\array_walk(
 			$this->services,
 			function ($class) {
-				// Load services classes but not in the WP-CLI env.
+				// Load services classes but not in the WP-CLI env, unless they have the ShouldLoadInCliContext attr.
 				if (!\defined('WP_CLI') && $class instanceof ServiceInterface) {
 					$class->register();
 				}
 
-				// Load services CLI classes only in WP-CLI env.
-				if (\defined('WP_CLI') && $class instanceof ServiceCliInterface) {
-					$class->register();
+				if (\defined('WP_CLI')) {
+					if ($class instanceof ServiceCliInterface) {
+						// Classes implementing ServiceCliInterface should be loaded only in CLI contexts.
+						$class->register();
+						return;
+					}
+
+					// Allow loading service classes in CLI contexts if it
+					// or a parent class has ShouldLoadInCliContext attribute.
+					$reflection = new ReflectionClass($class);
+					while ($reflection) {
+						if (\count($reflection->getAttributes(ShouldLoadInCliContext::class))) {
+							$class->register();
+							return;
+						}
+
+						$reflection = $reflection->getParentClass();
+					}
 				}
 			}
 		);