- Fix classmap tests (use get_class instead of instanceof)

- Add classmap to `MenuItemFactory`
- Don't return a menu when no `$params` are passed to `MenuFactory::from()`

lib/Factory/MenuFactory.php

@@ -49,11 +49,11 @@ public function from($params, array $args = []) : ?Menu {
 			}
 		}
 
-		if(!$menu) {
-			$menu = $this->from_nav_menu_terms($args);
+		if($menu) {
+			return $menu;
 		}
 
-		return $menu;
+		return $this->from_nav_menu_terms($args);
 	}
 
 	/**

lib/Factory/MenuItemFactory.php

@@ -38,7 +38,43 @@ protected function build(WP_Post $item, Menu $menu) : CoreInterface {
 		return $class::build($item, $menu);
 	}
 
-	protected function get_menuitem_class($item, $menu) : string {
+	protected function get_menuitem_class(WP_Post $item, Menu $menu) : string {
+		/**
+		 * Filters the class(es) used for different menu items.
+		 *
+		 * Read more about this in the documentation for [Menu Item Class Maps](https://timber.github.io/docs/v2/guides/class-maps/#the-menu-item-class-map).
+		 *
+		 * The default Menu Item Class Map will contain class names for locations that map to `Timber\MenuItem`.
+		 *
+		 * @since 2.0.0
+		 * @example
+		 * ```
+		 * add_filter( 'timber/menuitem/classmap', function( $classmap ) {
+		 *     $custom_classmap = [
+		 *         'primary'   => MenuItemFooter::class,
+		 *         'secondary' => MenuItemHeader::class,
+		 *     ];
+		 *
+		 *     return array_merge( $classmap, $custom_classmap );
+		 * } );
+		 * ```
+		 *
+		 * @param array $classmap The menu item class(es) to use. An associative array where the key is
+		 *                        the location and the value the name of the class to use for this
+		 *                        menu item or a callback that determines the class to use.
+		 */
+		$classmap = apply_filters( 'timber/menuitem/classmap', [] );
+
+		$class = $classmap[$menu->theme_location] ?? null;
+
+		// If class is a callable, call it to get the actual class name
+		if (is_callable($class)) {
+			$class = $class($item, $menu);
+		}
+
+		// Fallback on the default class
+		$class = $class ?? MenuItem::class;
+
 		/**
 		 * Filters the menu item class
 		 *
@@ -58,7 +94,7 @@ protected function get_menuitem_class($item, $menu) : string {
 		 * @param WP_Post $item The menu item.
 		 * @param Menu $menu The menu object.
 		 */
-		$class = apply_filters( 'timber/menuitem/class', MenuItem::class, $item, $menu );
+		$class = apply_filters( 'timber/menuitem/class', $class, $item, $menu );
 		return $class;
 	}
 }

lib/Menu.php

@@ -100,7 +100,7 @@ public static function build(?WP_Term $menu, $args = []) : ?self {
 		 *
 		 * @see wp_nav_menu()
 		 */
-		$defaults = array(
+		$defaults = [
 			'menu'                 => '',
 			'container'            => 'div',
 			'container_class'      => '',
@@ -119,11 +119,11 @@ public static function build(?WP_Term $menu, $args = []) : ?self {
 			'depth'                => 0,
 			'walker'               => '',
 			'theme_location'       => '',
-		);
+		];
 
 		$args = wp_parse_args( $args, $defaults );
 
-		if ( ! in_array( $args['item_spacing'], array( 'preserve', 'discard' ), true ) ) {
+		if ( ! in_array( $args['item_spacing'], ['preserve', 'discard'], true ) ) {
 			// Invalid value, fall back to default.
 			$args['item_spacing'] = $defaults['item_spacing'];
 		}
@@ -304,7 +304,7 @@ protected function init_pages_menu() {
 	 * @param array $items
 	 * @return MenuItem[]
 	 */
-	protected function convert_menu_items(array $menu_items) {
+	protected function convert_menu_items(array $menu_items) : array {
 		$menu_item_factory = new MenuItemFactory();
 		return array_map(function($item) use($menu_item_factory) : MenuItem {
 			return $menu_item_factory->from($item, $this);
@@ -418,7 +418,7 @@ public function get_items() {
 			return $this->items;
 		}
 
-		return array();
+		return [];
 	}
 
 	/**
@@ -484,7 +484,6 @@ public function current_top_level_item() {
 		return $this->current_item( 1 );
 	}
 
-
 	/**
 	 * Traverse an array of MenuItems in search of the current item.
 	 *
@@ -541,13 +540,13 @@ public function __toString() {
 
 		if ( $args->container ) {
 			/**
-				* Filters the list of HTML tags that are valid for use as menu containers.
-				*
-				* @since 3.0.0
-				*
-				* @param string[] $tags The acceptable HTML tags for use as menu containers.
-				*                       Default is array containing 'div' and 'nav'.
-				*/
+			* Filters the list of HTML tags that are valid for use as menu containers.
+			*
+			* @since 3.0.0
+			*
+			* @param string[] $tags The acceptable HTML tags for use as menu containers.
+			*                       Default is array containing 'div' and 'nav'.
+			*/
 			$allowed_tags = apply_filters( 'wp_nav_menu_container_allowedtags', array( 'div', 'nav' ) );
 
 			if ( is_string( $args->container ) && in_array( $args->container, $allowed_tags, true ) ) {

lib/MenuItem.php

@@ -14,7 +14,7 @@
  */
 class MenuItem extends CoreEntity {
 
-	protected WP_Post $core_object;
+	protected WP_Post $wp_object;
 
 	/**
 	 * @var string What does this class represent in WordPress terms?
@@ -100,11 +100,11 @@ public static function build( $data, ?Menu $menu = null ) : self {
 
 	/**
 	 * @internal
-	 * @param array|object $data
+	 * @param WP_Post $data
 	 * @param \Timber\Menu $menu The `Timber\Menu` object the menu item is associated with.
 	 */
-	protected function __construct( $data, $menu = null ) {
-		$this->core_object = $data;
+	protected function __construct( WP_Post $data, $menu = null ) {
+		$this->wp_object = $data;
 		$this->menu = $menu;
 
 		$data       = (object) $data;
@@ -306,7 +306,7 @@ public function import_classes( $data ) {
 		$this->classes = apply_filters(
 			'nav_menu_css_class',
 			$this->classes,
-			$this->core_object,
+			$this->wp_object,
 			$args,
 			0 // TODO: find the right depth
 		);
@@ -531,7 +531,7 @@ public function title() {
 			/**
 			 * @see Walker_Nav_Menu::start_el()
 			 */
-			$title = apply_filters( 'nav_menu_item_title', $this->__title, $this->core_object, $this->menu->args ? $this->menu->args : new \stdClass, $this->level );
+			$title = apply_filters( 'nav_menu_item_title', $this->__title, $this->wp_object, $this->menu->args ? $this->menu->args : new \stdClass, $this->level );
 			return $title;
 		}
 	}

tests/test-menu-factory.php

@@ -146,7 +146,6 @@ public function testFromWpTermObject() {
 			'taxonomy' => 'nav_menu',
 		]);
 
-
 		$factory = new MenuFactory();
 		$term    = get_term($id);
 
@@ -165,7 +164,7 @@ public function testMenuClassFilter() {
 			return MyMenu::class;
 		});
 
-		$this->assertInstanceOf(MyMenu::class, $factory->from($id));
+		$this->assertTrue(MyMenu::class === get_class($factory->from($id)));
 	}
 
 	public function testMenuClassMapFilter() {
@@ -188,6 +187,6 @@ public function testMenuClassMapFilter() {
 			];
 		});
 
-		$this->assertInstanceOf(MyMenu::class, $factory->from($id));
+		$this->assertTrue(MyMenu::class === get_class($factory->from($id)));
 	}
 }

tests/test-menu-item-factory.php

@@ -79,7 +79,7 @@ public function testMenuFromPost() {
 		$this->assertInstanceOf(MenuItem::class, $factory->from($post, $menu));
 	}
 
-	public function testMenuItemClassmap() {
+	public function testMenuItemClass() {
 		// Destructure the result into the menu WP_Term instance
 		// and the item_ids
 		[
@@ -120,8 +120,60 @@ public function testMenuItemClassmap() {
 			return $class;
 		}, 10, 3);
 
-		$this->assertInstanceOf(MenuItem::class,   $factory->from($one, $menu));
-		$this->assertInstanceOf(MyMenuItem::class, $factory->from($two, $menu));
-		$this->assertInstanceOf(MenuItem::class,   $factory->from($three, $menu));
+		$this->assertTrue(MenuItem::class === get_class($factory->from($one, $menu)));
+		$this->assertTrue(MyMenuItem::class === get_class($factory->from($two, $menu)));
+		$this->assertTrue(MenuItem::class === get_class($factory->from($three, $menu)));
 	}
+
+	public function testMenuItemClassmap() {
+		// Destructure the result into the menu WP_Term instance
+		// and the item_ids
+		[
+			'term' => $menu_term,
+			'item_ids' => [$one, $two, $three],
+		] = $this->create_menu_from_posts([
+			[
+				'post_title' => 'Page One',
+				'post_status' => 'publish',
+				'post_name' => 'page-one',
+				'post_type' => 'page',
+				'menu_order' => 1,
+			],
+			[
+				'post_title' => 'Page Two',
+				'post_status' => 'publish',
+				'post_name' => 'page-two',
+				'post_type' => 'page',
+				'menu_order' => 2,
+			],
+			[
+				'post_title' => 'Page Three',
+				'post_status' => 'publish',
+				'post_name' => 'page-three',
+				'post_type' => 'page',
+				'menu_order' => 3,
+			],
+		]);
+
+		register_nav_menu('custom', 'Custom nav location');
+		set_theme_mod('nav_menu_locations', [
+			'custom' => $menu_term['term_id'],
+		]);
+
+		$menu = Timber::get_menu($menu_term['term_id']);
+		$factory = new MenuItemFactory();
+
+		$this->add_filter_temporarily('timber/menuitem/classmap', function() {
+			return [
+				'custom' => MyMenuItem::class,
+			];
+		});
+
+		// Don't use instanceOf, it will return true whether it's MyMenuItem or MenuItem class
+		// and thus does not properly checks the classmap
+		$this->assertTrue(MyMenuItem::class === get_class($factory->from($one, $menu)));
+		$this->assertTrue(MyMenuItem::class === get_class($factory->from($two, $menu)));
+		$this->assertTrue(MyMenuItem::class === get_class($factory->from($three, $menu)));
+	}
+
 }

tests/test-timber-menu.php

@@ -931,7 +931,7 @@ function testMenuWalker() {
 		$args['echo'] = false;
 
 		$nav_menu_wp = wp_nav_menu($args);
-		// Prevents double ids to render
+		// Remove this filter that prevents to render duplicate ids
 		remove_filter( 'nav_menu_item_id', '_nav_menu_item_id_use_once', 10, 2 );
 		$nav_menu_timber = (string) $menu;