diff --git a/lib/src/models/node.dart b/lib/src/models/node.dart
index 16ac91a..acc8aec 100644
--- a/lib/src/models/node.dart
+++ b/lib/src/models/node.dart
@@ -33,7 +33,7 @@ class Node<T> {
   final T? data;
 
   /// The sub [Node]s of this object.
-  final List<Node> children;
+  final List<Node<T>> children;
 
   /// Force the node to be a parent so that node can show expander without
   /// having children node.
@@ -52,7 +52,7 @@ class Node<T> {
   /// Creates a [Node] from a string value. It generates a unique key.
   factory Node.fromLabel(String label) {
     String _key = Utilities.generateRandom();
-    return Node(
+    return Node<T>(
       key: '${_key}_$label',
       label: label,
     );
@@ -68,7 +68,7 @@ class Node<T> {
     String? _key = map['key'];
     String _label = map['label'];
     var _data = map['data'];
-    List<Node> _children = [];
+    List<Node<T>> _children = [];
     if (_key == null) {
       _key = Utilities.generateRandom();
     }
@@ -86,10 +86,10 @@ class Node<T> {
     if (map['children'] != null) {
       List<Map<String, dynamic>> _childrenMap = List.from(map['children']);
       _children = _childrenMap
-          .map((Map<String, dynamic> child) => Node.fromMap(child))
+          .map((Map<String, dynamic> child) => Node<T>.fromMap(child))
           .toList();
     }
-    return Node(
+    return Node<T>(
       key: '$_key',
       label: _label,
       data: _data,
@@ -101,16 +101,16 @@ class Node<T> {
 
   /// Creates a copy of this object but with the given fields
   /// replaced with the new values.
-  Node copyWith({
+  Node<T> copyWith({
     String? key,
     String? label,
-    List<Node>? children,
+    List<Node<T>>? children,
     bool? expanded,
     bool? parent,
     IconData? icon,
     T? data,
   }) =>
-      Node(
+      Node<T>(
         key: key ?? this.key,
         label: label ?? this.label,
         icon: icon ?? this.icon,
@@ -167,7 +167,7 @@ class Node<T> {
   bool operator ==(Object other) {
     if (identical(this, other)) return true;
     if (other.runtimeType != runtimeType) return false;
-    return other is Node &&
+    return other is Node<T> &&
         other.key == key &&
         other.label == label &&
         other.icon == icon &&
diff --git a/lib/src/tree_node.dart b/lib/src/tree_node.dart
index 2e4ab6a..2181641 100644
--- a/lib/src/tree_node.dart
+++ b/lib/src/tree_node.dart
@@ -19,17 +19,17 @@ const double _kBorderWidth = 0.75;
 /// __This class should not be used directly!__
 /// The [TreeView] and [TreeViewController] handlers the data and rendering
 /// of the nodes.
-class TreeNode extends StatefulWidget {
+class TreeNode<T> extends StatefulWidget {
   /// The node object used to display the widget state
-  final Node node;
+  final Node<T> node;
 
   const TreeNode({Key? key, required this.node}) : super(key: key);
 
   @override
-  _TreeNodeState createState() => _TreeNodeState();
+  _TreeNodeState<T> createState() => _TreeNodeState<T>();
 }
 
-class _TreeNodeState extends State<TreeNode>
+class _TreeNodeState<T> extends State<TreeNode<T>>
     with SingleTickerProviderStateMixin {
   static final Animatable<double> _easeInTween =
       CurveTween(curve: Curves.easeIn);
@@ -51,8 +51,10 @@ class _TreeNodeState extends State<TreeNode>
   @override
   void didChangeDependencies() {
     super.didChangeDependencies();
-    TreeView? _treeView = TreeView.of(context);
-    _controller.duration = _treeView!.theme.expandSpeed;
+    TreeView<T>? _treeView = TreeView.of<T>(context);
+    if (_treeView != null) {
+      _controller.duration = _treeView.theme.expandSpeed;
+    }
   }
 
   @override
@@ -62,7 +64,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   @override
-  void didUpdateWidget(TreeNode oldWidget) {
+  void didUpdateWidget(TreeNode<T> oldWidget) {
     if (widget.node.expanded != oldWidget.node.expanded) {
       setState(() {
         _isExpanded = widget.node.expanded;
@@ -82,7 +84,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   void _handleExpand() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     setState(() {
       _isExpanded = !_isExpanded;
@@ -100,7 +102,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   void _handleTap() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     if (_treeView!.onNodeTap != null) {
       _treeView.onNodeTap!(widget.node.key);
@@ -108,7 +110,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   void _handleDoubleTap() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     if (_treeView!.onNodeDoubleTap != null) {
       _treeView.onNodeDoubleTap!(widget.node.key);
@@ -116,7 +118,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   Widget _buildNodeExpander() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     TreeViewTheme _theme = _treeView!.theme;
     return widget.node.isParent
@@ -132,7 +134,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   Widget _buildNodeIcon() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     TreeViewTheme _theme = _treeView!.theme;
     bool isSelected = _treeView.controller.selectedKey != null &&
@@ -154,7 +156,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   Widget _buildNodeLabel() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     TreeViewTheme _theme = _treeView!.theme;
     bool isSelected = _treeView.controller.selectedKey != null &&
@@ -198,7 +200,7 @@ class _TreeNodeState extends State<TreeNode>
   }
 
   Widget _buildNodeWidget() {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     TreeViewTheme _theme = _treeView!.theme;
     bool isSelected = _treeView.controller.selectedKey != null &&
@@ -265,7 +267,7 @@ class _TreeNodeState extends State<TreeNode>
 
   @override
   Widget build(BuildContext context) {
-    TreeView? _treeView = TreeView.of(context);
+    TreeView<T>? _treeView = TreeView.of<T>(context);
     assert(_treeView != null, 'TreeView must exist in context');
     final bool closed =
         (!_isExpanded || !widget.node.expanded) && _controller.isDismissed;
@@ -295,8 +297,8 @@ class _TreeNodeState extends State<TreeNode>
                             _treeView.theme.iconTheme.size!),
                     child: Column(
                         mainAxisSize: MainAxisSize.min,
-                        children: widget.node.children.map((Node node) {
-                          return TreeNode(node: node);
+                        children: widget.node.children.map((Node<T> node) {
+                          return TreeNode<T>(node: node);
                         }).toList()),
                   ),
           )
diff --git a/lib/src/tree_view.dart b/lib/src/tree_view.dart
index cfa9d90..057bc1d 100644
--- a/lib/src/tree_view.dart
+++ b/lib/src/tree_view.dart
@@ -25,15 +25,15 @@ import 'models/node.dart';
 ///   theme: treeViewTheme
 /// ),
 /// ```
-class TreeView extends InheritedWidget {
+class TreeView<T> extends InheritedWidget {
   /// The controller for the [TreeView]. It manages the data and selected key.
-  final TreeViewController controller;
+  final TreeViewController<T> controller;
 
   /// The tap handler for a node. Passes the node key.
   final Function(String)? onNodeTap;
 
   /// Custom builder for nodes. Parameters are the build context and tree node.
-  final Widget Function(BuildContext, Node)? nodeBuilder;
+  final Widget Function(BuildContext, Node<T>)? nodeBuilder;
 
   /// The double tap handler for a node. Passes the node key.
   final Function(String)? onNodeDoubleTap;
@@ -92,7 +92,7 @@ class TreeView extends InheritedWidget {
   })  : this.theme = theme ?? const TreeViewTheme(),
         super(
           key: key,
-          child: _TreeViewData(
+          child: _TreeViewData<T>(
             controller,
             shrinkWrap: shrinkWrap,
             primary: primary,
@@ -100,7 +100,7 @@ class TreeView extends InheritedWidget {
           ),
         );
 
-  static TreeView? of(BuildContext context) =>
+  static TreeView<T>? of<T>(BuildContext context) =>
       context.dependOnInheritedWidgetOfExactType(aspect: TreeView);
 
   @override
@@ -114,8 +114,8 @@ class TreeView extends InheritedWidget {
   }
 }
 
-class _TreeViewData extends StatelessWidget {
-  final TreeViewController _controller;
+class _TreeViewData<T> extends StatelessWidget {
+  final TreeViewController<T> _controller;
   final bool? shrinkWrap;
   final bool? primary;
   final ScrollPhysics? physics;
@@ -133,8 +133,8 @@ class _TreeViewData extends StatelessWidget {
         primary: primary,
         physics: physics,
         padding: EdgeInsets.zero,
-        children: _controller.children.map((Node node) {
-          return TreeNode(node: node);
+        children: _controller.children.map<TreeNode<T>>((Node<T> node) {
+          return TreeNode<T>(node: node);
         }).toList(),
       ),
     );
diff --git a/lib/src/tree_view_controller.dart b/lib/src/tree_view_controller.dart
index 6f1fa9c..f3818b9 100644
--- a/lib/src/tree_view_controller.dart
+++ b/lib/src/tree_view_controller.dart
@@ -30,9 +30,9 @@ enum InsertMode {
 /// List<Node> newChildren = controller.updateNode(node.key, updatedNode);
 /// controller = TreeViewController(children: newChildren);
 /// ```
-class TreeViewController {
+class TreeViewController<T> {
   /// The data for the [TreeView].
-  final List<Node> children;
+  final List<Node<T>> children;
 
   /// The key of the select node in the [TreeView].
   final String? selectedKey;
@@ -44,8 +44,9 @@ class TreeViewController {
 
   /// Creates a copy of this controller but with the given fields
   /// replaced with the new values.
-  TreeViewController copyWith({List<Node>? children, String? selectedKey}) {
-    return TreeViewController(
+  TreeViewController<T> copyWith(
+      {List<Node<T>>? children, String? selectedKey}) {
+    return TreeViewController<T>(
       children: children ?? this.children,
       selectedKey: selectedKey ?? this.selectedKey,
     );
@@ -74,8 +75,8 @@ class TreeViewController {
   /// });
   /// ```
   TreeViewController loadMap({List<Map<String, dynamic>> list: const []}) {
-    List<Node> treeData =
-        list.map((Map<String, dynamic> item) => Node.fromMap(item)).toList();
+    List<Node<T>> treeData =
+        list.map((Map<String, dynamic> item) => Node<T>.fromMap(item)).toList();
     return TreeViewController(
       children: treeData,
       selectedKey: this.selectedKey,
@@ -96,12 +97,12 @@ class TreeViewController {
   /// ```
   TreeViewController withAddNode(
     String key,
-    Node newNode, {
-    Node? parent,
+    Node<T> newNode, {
+    Node<T>? parent,
     int? index,
     InsertMode mode: InsertMode.append,
   }) {
-    List<Node> _data =
+    List<Node<T>> _data =
         addNode(key, newNode, parent: parent, mode: mode, index: index);
     return TreeViewController(
       children: _data,
@@ -121,8 +122,9 @@ class TreeViewController {
   ///   controller = controller.withUpdateNode(key, newNode);
   /// });
   /// ```
-  TreeViewController withUpdateNode(String key, Node newNode, {Node? parent}) {
-    List<Node> _data = updateNode(key, newNode, parent: parent);
+  TreeViewController<T> withUpdateNode(String key, Node<T> newNode,
+      {Node<T>? parent}) {
+    List<Node<T>> _data = updateNode(key, newNode, parent: parent);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -141,8 +143,8 @@ class TreeViewController {
   ///   controller = controller.withDeleteNode(key);
   /// });
   /// ```
-  TreeViewController withDeleteNode(String key, {Node? parent}) {
-    List<Node> _data = deleteNode(key, parent: parent);
+  TreeViewController<T> withDeleteNode(String key, {Node<T>? parent}) {
+    List<Node<T>> _data = deleteNode(key, parent: parent);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -161,8 +163,8 @@ class TreeViewController {
   ///   controller = controller.withToggleNode(key, newNode);
   /// });
   /// ```
-  TreeViewController withToggleNode(String key, {Node? parent}) {
-    List<Node> _data = toggleNode(key, parent: parent);
+  TreeViewController withToggleNode(String key, {Node<T>? parent}) {
+    List<Node<T>> _data = toggleNode(key, parent: parent);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -182,7 +184,7 @@ class TreeViewController {
   /// });
   /// ```
   TreeViewController withExpandToNode(String key) {
-    List<Node> _data = expandToNode(key);
+    List<Node<T>> _data = expandToNode(key);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -202,7 +204,7 @@ class TreeViewController {
   /// });
   /// ```
   TreeViewController withCollapseToNode(String key) {
-    List<Node> _data = collapseToNode(key);
+    List<Node<T>> _data = collapseToNode(key);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -221,8 +223,8 @@ class TreeViewController {
   ///   controller = controller.withExpandAll();
   /// });
   /// ```
-  TreeViewController withExpandAll({Node? parent}) {
-    List<Node> _data = expandAll(parent: parent);
+  TreeViewController withExpandAll({Node<T>? parent}) {
+    List<Node<T>> _data = expandAll(parent: parent);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -241,8 +243,8 @@ class TreeViewController {
   ///   controller = controller.withCollapseAll();
   /// });
   /// ```
-  TreeViewController withCollapseAll({Node? parent}) {
-    List<Node> _data = collapseAll(parent: parent);
+  TreeViewController withCollapseAll({Node<T>? parent}) {
+    List<Node<T>> _data = collapseAll(parent: parent);
     return TreeViewController(
       children: _data,
       selectedKey: this.selectedKey,
@@ -250,12 +252,12 @@ class TreeViewController {
   }
 
   /// Gets the node that has a key value equal to the specified key.
-  Node? getNode(String key, {Node? parent}) {
-    Node? _found;
-    List<Node> _children = parent == null ? this.children : parent.children;
+  Node<T>? getNode(String key, {Node<T>? parent}) {
+    Node<T>? _found;
+    List<Node<T>> _children = parent == null ? this.children : parent.children;
     Iterator iter = _children.iterator;
     while (iter.moveNext()) {
-      Node child = iter.current;
+      Node<T> child = iter.current;
       if (child.key == key) {
         _found = child;
         break;
@@ -272,12 +274,12 @@ class TreeViewController {
   }
 
   /// Expands all node that are children of the parent node parameter. If no parent is passed, uses the root node as the parent.
-  List<Node> expandAll({Node? parent}) {
-    List<Node> _children = [];
+  List<Node<T>> expandAll({Node<T>? parent}) {
+    List<Node<T>> _children = [];
     Iterator iter =
         parent == null ? this.children.iterator : parent.children.iterator;
     while (iter.moveNext()) {
-      Node child = iter.current;
+      Node<T> child = iter.current;
       if (child.isParent) {
         _children.add(child.copyWith(
           expanded: true,
@@ -291,12 +293,12 @@ class TreeViewController {
   }
 
   /// Collapses all node that are children of the parent node parameter. If no parent is passed, uses the root node as the parent.
-  List<Node> collapseAll({Node? parent}) {
-    List<Node> _children = [];
+  List<Node<T>> collapseAll({Node<T>? parent}) {
+    List<Node<T>> _children = [];
     Iterator iter =
         parent == null ? this.children.iterator : parent.children.iterator;
     while (iter.moveNext()) {
-      Node child = iter.current;
+      Node<T> child = iter.current;
       if (child.isParent) {
         _children.add(child.copyWith(
           expanded: false,
@@ -310,12 +312,12 @@ class TreeViewController {
   }
 
   /// Gets the parent of the node identified by specified key.
-  Node? getParent(String key, {Node? parent}) {
-    Node? _found;
-    List<Node> _children = parent == null ? this.children : parent.children;
+  Node<T>? getParent(String key, {Node<T>? parent}) {
+    Node<T>? _found;
+    List<Node<T>> _children = parent == null ? this.children : parent.children;
     Iterator iter = _children.iterator;
     while (iter.moveNext()) {
-      Node child = iter.current;
+      Node<T> child = iter.current;
       if (child.key == key) {
         _found = parent ?? child;
         break;
@@ -333,23 +335,23 @@ class TreeViewController {
 
   /// Expands a node and all of the node's ancestors so that the node is
   /// visible without the need to manually expand each node.
-  List<Node> expandToNode(String key) {
+  List<Node<T>> expandToNode(String key) {
     List<String> _ancestors = [];
     String _currentKey = key;
 
     _ancestors.add(_currentKey);
 
-    Node? _parent = this.getParent(_currentKey);
+    Node<T>? _parent = this.getParent(_currentKey);
     if (_parent != null) {
       while (_parent!.key != _currentKey) {
         _currentKey = _parent.key;
         _ancestors.add(_currentKey);
         _parent = this.getParent(_currentKey);
       }
-      TreeViewController _this = this;
+      TreeViewController<T> _this = this;
       _ancestors.forEach((String k) {
-        Node _node = _this.getNode(k)!;
-        Node _updated = _node.copyWith(expanded: true);
+        Node<T> _node = _this.getNode(k)!;
+        Node<T> _updated = _node.copyWith(expanded: true);
         _this = _this.withUpdateNode(k, _updated);
       });
       return _this.children;
@@ -359,23 +361,23 @@ class TreeViewController {
 
   /// Collapses a node and all of the node's ancestors without the need to
   /// manually collapse each node.
-  List<Node> collapseToNode(String key) {
+  List<Node<T>> collapseToNode(String key) {
     List<String> _ancestors = [];
     String _currentKey = key;
 
     _ancestors.add(_currentKey);
 
-    Node? _parent = this.getParent(_currentKey);
+    Node<T>? _parent = this.getParent(_currentKey);
     if (_parent != null) {
       while (_parent!.key != _currentKey) {
         _currentKey = _parent.key;
         _ancestors.add(_currentKey);
         _parent = this.getParent(_currentKey);
       }
-      TreeViewController _this = this;
+      TreeViewController<T> _this = this;
       _ancestors.forEach((String k) {
-        Node _node = _this.getNode(k)!;
-        Node _updated = _node.copyWith(expanded: false);
+        Node<T> _node = _this.getNode(k)!;
+        Node<T> _updated = _node.copyWith(expanded: false);
         _this = _this.withUpdateNode(k, _updated);
       });
       return _this.children;
@@ -387,17 +389,17 @@ class TreeViewController {
   /// accepts an [InsertMode] and index. If no [InsertMode] is specified,
   /// it appends the new node as a child at the end. This method returns
   /// a new list with the added node.
-  List<Node> addNode(
+  List<Node<T>> addNode(
     String key,
-    Node newNode, {
-    Node? parent,
+    Node<T> newNode, {
+    Node<T>? parent,
     int? index,
     InsertMode mode: InsertMode.append,
   }) {
-    List<Node> _children = parent == null ? this.children : parent.children;
-    return _children.map((Node child) {
+    List<Node<T>> _children = parent == null ? this.children : parent.children;
+    return _children.map((Node<T> child) {
       if (child.key == key) {
-        List<Node> _children = child.children.toList(growable: true);
+        List<Node<T>> _children = child.children.toList(growable: true);
         if (mode == InsertMode.prepend) {
           _children.insert(0, newNode);
         } else if (mode == InsertMode.insert) {
@@ -422,9 +424,9 @@ class TreeViewController {
 
   /// Updates an existing node identified by specified key. This method
   /// returns a new list with the updated node.
-  List<Node> updateNode(String key, Node newNode, {Node? parent}) {
-    List<Node> _children = parent == null ? this.children : parent.children;
-    return _children.map((Node child) {
+  List<Node<T>> updateNode(String key, Node<T> newNode, {Node<T>? parent}) {
+    List<Node<T>> _children = parent == null ? this.children : parent.children;
+    return _children.map((Node<T> child) {
       if (child.key == key) {
         return newNode;
       } else {
@@ -444,19 +446,19 @@ class TreeViewController {
 
   /// Toggles an existing node identified by specified key. This method
   /// returns a new list with the specified node toggled.
-  List<Node> toggleNode(String key, {Node? parent}) {
-    Node? _node = getNode(key, parent: parent);
+  List<Node<T>> toggleNode(String key, {Node<T>? parent}) {
+    Node<T>? _node = getNode(key, parent: parent);
     return updateNode(key, _node!.copyWith(expanded: !_node.expanded));
   }
 
   /// Deletes an existing node identified by specified key. This method
   /// returns a new list with the specified node removed.
-  List<Node> deleteNode(String key, {Node? parent}) {
-    List<Node> _children = parent == null ? this.children : parent.children;
-    List<Node> _filteredChildren = [];
+  List<Node<T>> deleteNode(String key, {Node<T>? parent}) {
+    List<Node<T>> _children = parent == null ? this.children : parent.children;
+    List<Node<T>> _filteredChildren = [];
     Iterator iter = _children.iterator;
     while (iter.moveNext()) {
-      Node child = iter.current;
+      Node<T> child = iter.current;
       if (child.key != key) {
         if (child.isParent) {
           _filteredChildren.add(child.copyWith(
@@ -471,13 +473,13 @@ class TreeViewController {
   }
 
   /// Get the current selected node. Returns null if there is no selectedKey
-  Node? get selectedNode {
+  Node<T>? get selectedNode {
     return this.selectedKey!.isEmpty ? null : getNode(this.selectedKey!);
   }
 
   /// Map representation of this object
   List<Map<String, dynamic>> get asMap {
-    return children.map((Node child) => child.asMap).toList();
+    return children.map((Node<T> child) => child.asMap).toList();
   }
 
   @override
diff --git a/pubspec.yaml b/pubspec.yaml
index 6beff3c..86e0044 100644
--- a/pubspec.yaml
+++ b/pubspec.yaml
@@ -1,6 +1,6 @@
 name: flutter_treeview
 description: A tree widget for Flutter that can be used to display nested, hierarchical data. It includes a number of features like styling labels, icons, and import and export utilities.
-version: 1.0.3+2
+version: 1.0.3+16
 homepage: https://bitbucket.org/kevinandre/flutter_treeview/src/master/
 repository: https://bitbucket.org/kevinandre/flutter_treeview/src/master/
 issue_tracker: https://bitbucket.org/kevinandre/flutter_treeview/issues