diff --git a/src/index.jsx b/src/index.tsx
similarity index 54%
rename from src/index.jsx
rename to src/index.tsx
index 981d67a..37d65eb 100644
--- a/src/index.jsx
+++ b/src/index.tsx
@@ -1,20 +1,45 @@
 import React, { createContext, useContext } from 'react'
-import PropTypes from 'prop-types'
 
-export const DEFAULT_TREE_CONTEXT = {
-  getId: (node) => node.id,
-  getChildNodes: (node) => node.childNodes,
-  renderLeaf: ({ id }, { level }) => (
+export type Node = {
+  id: string | number
+  childNodes?: Node[]
+  [key: string]: any
+}
+
+export type TreeContext = {
+  getChildNodes: (node: Node) => Node[]
+  renderBranch: (
+    node: Node,
+    opts: {
+      level: number
+      index: number
+      children: React.ReactNode
+      nodes: Node[]
+    }
+  ) => JSX.Element
+  renderLeaf: (
+    node: Node,
+    opts: {
+      level: number
+      index: number
+      nodes: Node[]
+    }
+  ) => JSX.Element
+}
+
+export const DEFAULT_TREE_CONTEXT: TreeContext = {
+  getChildNodes: (node) => node.childNodes || [],
+  renderLeaf: (node, { level }) => (
     <div className='TreeNode TreeNode--Leaf'>
       <div className='TreeNode__Label'>
-        {id} ({level})
+        {node.id} ({level})
       </div>
     </div>
   ),
-  renderBranch: ({ id }, { children, level }) => (
+  renderBranch: (node, { children, level }) => (
     <div className='TreeNode TreeNode--Branch'>
       <div className='TreeNode__Label'>
-        {id} ({level})
+        {node.id} ({level})
       </div>
       <div className='TreeNode__ChildNodeContainer'>{children}</div>
     </div>
@@ -23,14 +48,18 @@ export const DEFAULT_TREE_CONTEXT = {
 
 export const TreeContext = createContext(DEFAULT_TREE_CONTEXT)
 
-export const TreeNodeList = ({ nodes, level = 0 }) => {
-  const { getId } = useContext(TreeContext)
-
+export const TreeNodeList = ({
+  nodes,
+  level = 0,
+}: {
+  nodes: Node[]
+  level: number
+}) => {
   return (
     <React.Fragment>
       {nodes.map((node, index, nodes) => (
         <TreeNode
-          key={getId(node)}
+          key={node.id}
           node={node}
           level={level}
           index={index}
@@ -41,11 +70,18 @@ export const TreeNodeList = ({ nodes, level = 0 }) => {
   )
 }
 TreeNodeList.displayName = 'TreeNodeList'
-TreeNodeList.propTypes = {
-  nodes: PropTypes.array.isRequired,
-}
 
-export const TreeNode = ({ node, level = 0, index, nodes }) => {
+export const TreeNode = ({
+  node,
+  level = 0,
+  index,
+  nodes,
+}: {
+  node: Node
+  level?: number
+  index: number
+  nodes: Node[]
+}) => {
   const { getChildNodes, renderLeaf, renderBranch } = useContext(TreeContext)
   const childNodes = getChildNodes(node)
 
@@ -63,6 +99,3 @@ export const TreeNode = ({ node, level = 0, index, nodes }) => {
       })
 }
 TreeNode.displayName = 'TreeNode'
-TreeNode.propTypes = {
-  node: PropTypes.object.isRequired,
-}